Cannot export PyTorch model to ONNX

Question

I am trying to convert a pre-trained torch model to ONNX, but recive the following error:

RuntimeError: step!=1 is currently not supported

I'm trying this on a pre-trained colorization model: https://github.com/richzhang/colorization

Here is the code I ran in 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)

I appreciate any help :)

UPDATE 1: @Proko suggestion solved the ONNX export issue. Now I have a new possibly related problem when I try to convert the ONNX to TensorRT. I get the following error:

[TensorRT] ERROR: Network must have at least one output

Here is the code I used:

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)

I tried to force the build_engine function to use the output of the network by:

network.mark_output(network.get_layer(network.num_layers-1).get_output(0))

but it did not work. I appropriate any help!

you will not be able to export this model at this time. There are slices with step different than 1 which are simply not supported right now by `torch.onnx` . Maybe rewriting the model to use something different than `n`-step slice (but of course gives the same result) might help you — Proko, Dec 29 '20 at 19:59

Proko · Accepted Answer · 2020-12-29T22:25:46.007

Like I have mentioned in a comment, this is because slicing in torch.onnx supports only step = 1 but there are 2-step slicing in the model:

self.model2(conv1_2[:,:,::2,::2])

Your only option as for now is to rewrite slicing to be some other ops. You can do it by using range and reshape to obtain proper indices. Consider the following function "step-less-arange" (I hope it is generic enough for anyone with similar problem):

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]

usage:

>> sla(11, 3)
tensor([0, 3, 6, 9])

Now you can replace every slice like this:

conv2_2 = self.model2(conv1_2[:,:,self.sla(conv1_2.shape[2], 2),:][:,:,:, self.sla(conv1_2.shape[3], 2)])

NOTE: you should optimize it. Indices are calculated for every call so it might be wise to pre-compute it.

I have tested it with my fork of the repo and I was able to save the model:

https://github.com/prokotg/colorization

Thanks a lot! It worked for exporting the model. Now I have a possibly related issue with converting the ONNX to TRT. I'll appreciate if you take a look at mt edited question and let me know if you have any idea how to solve it :) — Darkoob, Dec 29 '20 at 23:45
@Darkoob sorry, I don't have an access to tensorrt right now (don't have a gpu). I might get back to you once I have access to a machine with gpu — Proko, Dec 30 '20 at 12:08

score 1 · Answer 2 · answered Aug 18 '21 at 17:14

What works for me was to add the opset_version=11 on torch.onnx.export

First I had tried use opset_version=10, but the API suggest 11 so it works.

So your function should be:

torch.onnx.export(model, dummy_input, "test_converted_model.onnx", verbose=True,opset_version=11, input_names=input_names, output_names=output_names)

Cannot export PyTorch model to ONNX

2 Answers2