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

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

0 голосов
/ 18 января 2019

Я пытаюсь преобразовать модель 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

1 Ответ

0 голосов
/ 24 января 2019

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

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

...