我正在尝试使用以下代码将Pytorch模型转换为ONNX:
import onnx
import torch
import onnxruntime
import numpy as np
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
batch_size = 1
model = torch.jit.load('../saved_weights/autotissue_mobilenetv3_lrassp.pt')
x = torch.randn(1, 3, 224, 224, requires_grad=True)
x = x.to(device)
torch_out = model(x)
torch_out = torch_out['out']
torch.onnx.export(model, x, "../saved_weights/autotissue_mobilenetV3.onnx",
export_params=True,
opset_version=12,
do_constant_folding=True,
input_names=['input'],
output_names=['out'],
dynamic_axes={ 'input': {0 : 'batch_size'},
'output': {0 : 'batch_size'}})
但是,我得到以下错误:
raise errors.UnsupportedOperatorError(
torch.onnx.errors.UnsupportedOperatorError: Exporting the operator
'aten::dict' to ONNX opset version 12 is not supported.
我不知道为什么,但是会不会是模型输出是字典的形式呢?
1条答案
按热度按时间v7pvogib1#
参考Pytorch文档,您可以尝试:
1.更改模型以不使用该操作符(此处为字典)。
1.创建一个符号函数来转换运算符并将其注册为自定义符号函数。
1.为PyTorch做贡献,将相同的符号函数添加到torch.onnx本身。
我的建议是将输出类型更改为Tensor。
此外,您还可以访问here并检查支持的操作符。