无法将 PyTorch 模型导出到 ONNX
Posted
技术标签:
【中文标题】无法将 PyTorch 模型导出到 ONNX【英文标题】:Cannot export PyTorch model to ONNX 【发布时间】:2021-04-06 07:56:28 【问题描述】:我正在尝试将预训练的 Torch 模型转换为 ONNX,但收到以下错误:
RuntimeError: step!=1 is currently not supported
我正在预训练的着色模型上尝试这个:https://github.com/richzhang/colorization
这是我在 Google Colab 中运行的代码:
!git clone https://github.com/richzhang/colorization.git
cd colorization/
import colorizers
model = colorizer_siggraph17 = colorizers.siggraph17(pretrained=True).eval()
input_names = [ "input" ]
output_names = [ "output" ]
dummy_input = torch.randn(1, 1, 256, 256, device='cpu')
torch.onnx.export(model, dummy_input, "test_converted_model.onnx", verbose=True,
input_names=input_names, output_names=output_names)
感谢您的帮助:)
更新 1: @Proko 建议解决了 ONNX 导出问题。现在,当我尝试将 ONNX 转换为 TensorRT 时,我遇到了一个可能相关的新问题。我收到以下错误:
[TensorRT] ERROR: Network must have at least one output
这是我使用的代码:
import torch
import pycuda.driver as cuda
import pycuda.autoinit
import tensorrt as trt
import onnx
TRT_LOGGER = trt.Logger()
def build_engine(onnx_file_path):
# initialize TensorRT engine and parse ONNX model
builder = trt.Builder(TRT_LOGGER)
builder.max_workspace_size = 1 << 25
builder.max_batch_size = 1
if builder.platform_has_fast_fp16:
builder.fp16_mode = True
network = builder.create_network()
parser = trt.OnnxParser(network, TRT_LOGGER)
# parse ONNX
with open(onnx_file_path, 'rb') as model:
print('Beginning ONNX file parsing')
parser.parse(model.read())
print('Completed parsing of ONNX file')
# generate TensorRT engine optimized for the target platform
print('Building an engine...')
engine = builder.build_cuda_engine(network)
context = engine.create_execution_context()
print("Completed creating Engine")
return engine, context
ONNX_FILE_PATH = 'siggraph17.onnx' # Exported using the code above
engine,_ = build_engine(ONNX_FILE_PATH)
我试图通过以下方式强制 build_engine 函数使用网络的输出:
network.mark_output(network.get_layer(network.num_layers-1).get_output(0))
但它不起作用。 我需要任何帮助!
【问题讨论】:
此时您将无法导出此模型。torch.onnx
现在根本不支持步长不同于 1 的切片。也许重写模型以使用与n
-step 切片不同的东西(但当然给出相同的结果)可能会对您有所帮助
【参考方案1】:
就像我在评论中提到的,这是因为torch.onnx
中的切片仅支持step = 1
,但模型中有两步切片:
self.model2(conv1_2[:,:,::2,::2])
目前您唯一的选择是将切片重写为其他一些操作。您可以通过使用 range 和 reshape 来获得适当的索引。考虑以下函数“step-less-arange”(我希望它对任何有类似问题的人来说足够通用):
def sla(x, step):
diff = x % step
x += (diff > 0)*(step - diff) # add length to be able to reshape properly
return torch.arange(x).reshape((-1, step))[:, 0]
用法:
>> sla(11, 3)
tensor([0, 3, 6, 9])
现在您可以像这样替换每个切片:
conv2_2 = self.model2(conv1_2[:,:,self.sla(conv1_2.shape[2], 2),:][:,:,:, self.sla(conv1_2.shape[3], 2)])
注意:您应该优化它。每次调用都会计算索引,因此最好预先计算它。
我已经使用我的 repo 分支对其进行了测试,并且我能够保存模型:
https://github.com/prokotg/colorization
【讨论】:
非常感谢!它适用于导出模型。现在我有一个可能与将 ONNX 转换为 TRT 相关的问题。如果您查看 mt 编辑的问题,我将不胜感激,如果您知道如何解决它,请告诉我:) @Darkoob 抱歉,我现在无法访问 tensorrt(没有 gpu)。一旦我可以使用带有 gpu 的机器,我可能会回复你【参考方案2】:对我有用的是在 torch.onnx.export 上添加 opset_version=11
首先我尝试使用 opset_version=10,但 API 建议使用 11,所以它可以工作。
所以你的功能应该是:
torch.onnx.export(model, dummy_input, "test_converted_model.onnx", verbose=True,opset_version=11, 输入名称=输入名称,输出名称=输出名称)
【讨论】:
以上是关于无法将 PyTorch 模型导出到 ONNX的主要内容,如果未能解决你的问题,请参考以下文章
PyTorch 1.0 中文官方教程:使用ONNX将模型从PyTorch传输到Caffe2和移动端
Pytorch 到 ONNX 导出功能失败并导致遗留功能错误
将 PyTorch 模型与 CoreML 一起使用时输入尺寸重塑