1

我正在尝试将 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
4

1 回答 1

1

你得到什么输出？如文档中所述，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 通道图像作为输入，对吧？

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

于 2019-01-24T10:12:46.303 回答