Блог

Я пытаюсь преобразовать предварительно подготовленную модель 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 с тем, что поддерживается только нарезка step = 1 , но в модели есть 2-ступенчатая нарезка:

 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)])
 

ПРИМЕЧАНИЕ: вы должны оптимизировать его. Индексы вычисляются для каждого вызова, поэтому было бы разумно предварительно вычислить его.

Я протестировал это с помощью своей вилки репо, и мне удалось сохранить модель:

https://github.com/prokotg/colorization

Что сработало для меня, так это добавить opset_version=11 в torch.onnx.export

Сначала я попытался использовать opset_version=10, но API предлагает 11, поэтому он работает.

Таким образом, ваша функция должна быть:

torch.onnx.export(модель, dummy_input, «test_converted_model.onnx», verbose= True,opset_version= 11, input_names=input_names, output_names=output_names)