Передача предварительно обученной модели pytorch в onnx

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

Какой выход вы получаете? Кажется, SuperResolution поддерживается с операторами экспорта в pytorch, как указано в документации

Вы уверены, что вход для вашей модели:

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-канального изображения в качестве входа, верно?

Кроме того, было бы полезно, если бы вы опубликовали вывод терминала, чтобы увидеть, где он выходит из строя. Удачи!