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



目标:使用此笔记本对albert-base-v2模型执行量化。

内核:conda_pytorch_p36


第1.2节中的输出;2.2显示:

  • 将香草BERTPyTorch转换为ONNX保持相同大小417.6 MB
  • 量化模型比vanilla BERT、PyTorch173.0 MB和ONNX104.8 MB更小

但是,当运行ALBert时:

  • PyTorch和ONNX型号尺寸不同
  • 量化模型尺寸比香草大

I认为这是与香草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

如果还有什么可以添加到帖子中的,请告诉我。

说明

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

关于这个现象,许多Git问题已被标记为已解决。

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


解决方案

章节";移除共享权重";在onnx_remove_shared_weights.ipynb.中

伪代码:

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)

两种解决方案的来源

最新更新