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!

Answer 1

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