将 PyTorch 转换为 ONNX 模型会增加 ALBert 的文件大小

Question

目标：使用此 Notebook 对 albert-base-v2 模型执行量化。

内核：conda_pytorch_p36.

第 1.2 和 2.2 节 中的输出表明：

将 vanilla BERT 从 PyTorch 转换为 ONNX 保持相同大小，417.6 MB.
量化模型比普通 BERT、PyTorch 173.0 MB 和 ONNX 104.8 MB.

更小

但是，当运行 ALBert:

PyTorch 和 ONNX 模型尺寸不同。
量化后的模型大小 比原始模型大 。

我认为这就是 ALBert 的两种量化方法与普通 ALBert 相比 模型性能较差的原因。

PyTorch:

Size (MB): 44.58906650543213 Size (MB): 22.373255729675293

ONNX:

ONNX full precision model size (MB): 341.64233207702637 ONNX quantized model size (MB): 85.53886985778809

为什么将 ALBert 从 PyTorch 导出到 ONNX 会增加模型大小，但不会增加 BERT？

如果我还有什么要补充的，请告诉我 post。

Answer 1

说明

ALBert 模型在层之间共享权重。 torch.onnx.export 将权重输出到不同的张量，这导致模型尺寸变大。

关于此现象的一些 Git 问题已标记为已解决。

最常见的解决方案是删除共享权重，即删除包含完全相同值的张量数组。

解决方案

onnx_remove_shared_weights.ipynb 中的“删除共享权重”部分。

Pseudo-code:

from onnxruntime.transformers.onnx_model import OnnxModel
model=onnx.load(path)
onnx_model=OnnxModel(model)
count = len(model.graph.initializer)
same = [-1] * count
for i in range(count - 1):
  if same[i] >= 0:
    continue
  for j in range(i+1, count):
     if has_same_value(model.graph.initializer[i], model.graph.initializer[j]):
    same[j] = i

for i in range(count):
   if same[i] >= 0:
        onnx_model.replace_input_of_all_nodes(model.graph.initializer[i].name, model.graph.initializer[same[i]].name)

onnx_model.update_graph()
onnx_model.save_model_to_file(output_path)

Source of both solutions

将 PyTorch 转换为 ONNX 模型会增加 ALBert 的文件大小

Converting PyTorch to ONNX model increases file size for ALBert

python

quantization

pytorch

onnx

onnxruntime

说明

解决方案