我正在尝试在 pytorch 中训练一个量化模型并将其转换为 ONNX。我在 pytorch_quantization 包的帮助下采用了量化感知训练技术。我使用以下代码将我的模型转换为 ONNX:
from pytorch_quantization import nn as quant_nn
from pytorch_quantization import calib
from pytorch_quantization.tensor_quant import QuantDescriptor
from pytorch_quantization import quant_modules
import onnxruntime
import torch
import torch.utils.data
from torch import nn
import torchvision
def export_onnx(model, onnx_filename, batch_onnx, per_channel_quantization):
model.eval()
quant_nn.TensorQuantizer.use_fb_fake_quant = True # We have to shift to pytorch's fake quant ops before exporting the model to ONNX
if per_channel_quantization:
opset_version = 13
else:
opset_version = 12
# Export ONNX for multiple batch sizes
print("Creating ONNX file: " + onnx_filename)
dummy_input = torch.randn(batch_onnx, 3, 224, 224, device='cuda') #TODO: switch input dims by model
input_names = ['input']
output_names = ['Linear[fc]'] ### ResNet34
dynamic_axes = {'input': {0: 'batch_size'}}
try:
torch.onnx.export(model, dummy_input, onnx_filename, input_names=input_names,
export_params=True, output_names=output_names, opset_version=opset_version,
verbose=True, enable_onnx_checker=False, do_constant_folding=True)
except ValueError:
warnings.warn(UserWarning("Per-channel quantization is not yet supported in Pytorch/ONNX RT (requires ONNX opset 13)"))
print("Failed to export to ONNX")
return False
return True
转换后,我收到以下警告:
warnings.warn("'enable_onnx_checker' 已弃用并忽略。它将在 " W0305 12:39:40.472136 140018114328384 tensor_quantizer.py:280] 中删除 使用 Pytorch 的原生实验性假量化。
/usr/local/lib/python3.8/dist-packages/pytorch_quantization/nn/modules/tensor_quantizer.py:285:TracerWarning:将张量转换为 Python 数字可能会导致跟踪不正确。我们无法记录 Python 值的数据流向,所以以后这个值会被当作一个常量来处理。这意味着跟踪可能不会推广到其他输入!
此外,精度对 ONNX 模型无效!
Accuracy summary:
+-----------+-------+
| Stage | Top1 |
+-----------+-------+
| Finetuned | 38.03 |
| ONNX | -1.00 |
+-----------+-------+
更多信息在这里:
pytorch 1.10.2+cu102
torchvision 0.11.3+cu102
TensorRT 8.2.3-1+cuda11.4
ONNX 1.11.0
ONNX Runtime 1.10.0
cuda 11.6
python 3.8
ONNX 转换有什么问题?