Question

我正在尝试将pytorch模型转换为ONNX，以便以后将其用于TensorRT。我遵循了以下教程https://pytorch.org/tutorials/advanced/super_resolution_with_caffe2.html，但是我的内核一直死掉。

这是我实现的代码。

 # Some standard imports
import io
import numpy as np

from torch import nn
import torch.onnx
from deepformer.nets.quicknat import quickNAT
param = {
    'num_channels': 64,
    'num_filters': 64,
    'kernel_h': 5,
    'kernel_w': 5,
    'kernel_c': 1,
    'stride_conv': 1,
    'pool': 2,
    'stride_pool': 2,
    'num_classes': 1,
    'padding': 'reflection'
}


net = quickNAT(param)
checkpoint_path = 'checkpoint_epoch36_loss0.78.t7'
checkpoints=torch.load(checkpoint_path)
map_location = lambda storage, loc: storage
if torch.cuda.is_available():
    map_location = None
net.load_state_dict(checkpoints['net'])
net.train(False)
# Input to the modelvcdfx  
x = torch.rand(1, 64, 256, 1600, requires_grad=True)

# Export the model
torch_out = torch.onnx._export(net,             # model being run
                               x,                       # model input (or a tuple for multiple inputs)
                               "quicknat.onnx", # where to save the model (can be a file or file-like object)
                               export_params=True)      # store the trained parameter weights inside the model file

Answer 1

您得到的输出是什么？似乎in the documentation

中的pytorch中的导出运算符支持SuperResolution。

您确定模型的输入是：

x = torch.rand(1, 64, 256, 1600, requires_grad=True)

这可能是您用于训练的变量，因为对于部署而言，您在一个或多个映像上运行网络，因此要导出到onnx的虚拟输入通常是：

dummy_input = torch.randn(1, 3, 720, 1280, device='cuda')

其中1是批处理大小，3是图像（RGB）的通道，然后是图像的大小，在这种情况下为720x1280。检查该输入，我想您输入的不是64通道图像吗？

此外，如果您发布终端输出以查看其失败位置，这将对您有所帮助。祝你好运！

将预训练的pytorch模型转移到onnx

1 个答案: