step!=1 is currently not supported

This error is raised when exporting a PyTorch model to onnx format. The cause of the error is when the torch model that is being converted to onnx format contains slices with step greater than 1.

What do we mean slices with step greater than 1?

A slice step is a number of jumps that tells what number of elements to jump over when indexing and selecting elements from a numpy array or torch tensor.

E.g

array_20 = np.arange(20)
selected = array_20[3:15:2]

print(selected)

From the above example the array_20 is assigned 20 consecutive integers. i.e

[0,1,2,3....19].

In the second line the array_20[3:15:2] is used to select the numbers from 2:15 in a way that each number is differed from its predecessor or successor by 2. This is indicated by the last index 2.

This is what we call slices with step=2.

If we want to select from all of the elements by setting a step=2, we use...

selected_all = array_20[::2] 

# here the start and end indice are removed and we are selecting from all the elements 

Now if our Torch model has an operation that resembles with above slicing and if the onnx version doesn't support that operation the mentioned error ...

RuntimeError("step!=1 is currently not supported")

is raised.

Steps to reproduce the error:

• Setup virtual environment

$ pip install --user pipenv
$ mkdir test_folder
$ cd test_folder
$ pipenv shell

• Install numpy

$ pipenv install numpy

• Install pytorch

Here we can install pytorch two ways, for cudaor for cpudepending on the machine being used for

test.

For Cuda:

This also depends on the version of the cuda installed in the OS

# CUDA 11.0
pipenv install torch==1.7.0+cu110 torchvision==0.8.0+cu110 torchaudio==0.7.0 -f https://download.pytorch.org/whl/torch_stable.html

# CUDA 10.2
pipenv install torch==1.7.0 torchvision==0.8.0 torchaudio==0.7.0

# CUDA 10.1
pipenv install torch==1.7.0+cu101 torchvision==0.8.0+cu101 torchaudio==0.7.0 -f https://download.pytorch.org/whl/torch_stable.html

# CUDA 9.2
pipenv install torch==1.7.0+cu92 torchvision==0.8.0+cu92 torchaudio==0.7.0 -f https://download.pytorch.org/whl/torch_stable.html

For CPU:

pipenv install torch==1.7.0+cpu torchvision==0.8.0+cpu torchaudio==0.7.0 -f https://download.pytorch.org/whl/torch_stable.html

• Run test code

import torch
import torch.nn as nn
import torch.nn.init as init
import numpy as np
import os
import tempfile

class CustomAdder(torch.nn.Module):
    def __init__(self):
        super().__init__()
    def forward(self, x):
        x_modified = x[::2] # Here our model has slice with step=2 (i.e step!=1)
        return 1+x_modified #
    def string(self):
        return f'y = 1+x'

device = 'cuda' if torch.cuda.is_available() else 'cpu'
# initialize the x arary
x = np.arange(10).reshape(10,1)
x_tensor = torch.from_numpy(x).float().to(device) #convert the np array to tensor and copy to device (cpu or cuda

# initialize our model and set it to eval(). Here since the model is a simple test model
# no training is done on it.

model = CustomAdder() 
model.eval()
# run forward functions of the model by calling the model object.
y = model(x_tensor)
print(y)

# Now let's try to export it an onnx format
with tempfile.TemporaryDirectory() as tmpdir:
    model_path = os.path.join(tmpdir,"customAdder.onnx")
    torch.onnx.export(model,x_tensor,model_path,
                        export_params=True,
                        do_constant_folding=True,
                        input_names = ['input'],
                        output_names = ['output'],
                        dynamic_axes={'input' : {0 : 'batch_size'},
                                      'output' : {0 : 'batch_size'}})
     
      #This raises the mentioned error

For more details on the onnx exporter check the tutorial on the official Pytorch website.

How to fix this Error:

There are two ways to fix this error

• By adding a custom slicer to avoid doing the automatic slicing

This means doing the slicing using a custom method to do it one by one.

import torch
import torch.nn as nn
import torch.nn.init as init
import numpy as np
import os
import tempfile

class CustomAdder(torch.nn.Module):
    def __init__(self):
        super().__init__()
    @staticmethod
    def sla(x, step):
        diff = x % step
        x += (diff > 0) * (step - diff)
        return torch.arange(x).reshape((-1, step))[:, 0]
    def forward(self, x):
        length = x_tensor.shape[0]
        #x_modified = x[::2] 
        x_modified = self.sla(length,2) # modifying x using the custom method
        return 1+x_modified #
    def string(self):
        return f'y = 1+x'

device = 'cuda' if torch.cuda.is_available() else 'cpu'
# initialize the x arary
x = np.arange(10).reshape(10,1)
x_tensor = torch.from_numpy(x).float().to(device) #convert the np array to tensor and copy to device (cpu or cuda

# initialize our model and set it to eval(). Here since the model is a simple test model
# no training is done on it.

model = CustomAdder() 
model.eval()
# run forward functions of the model by calling the model object.
y = model(x_tensor)
print(y)

# Now let's try to export it an onnx format
with tempfile.TemporaryDirectory() as tmpdir:
    model_path = os.path.join(tmpdir,"customAdder.onnx")  
    nnx.export(model,x_tensor,model_path,
                            export_params=True,
                            do_constant_folding=True,
                            input_names = ['input'],
                            output_names = ['output'],
                            dynamic_axes={'input' : {0 : 'batch_size'},
                                          'output' : {0 : 'batch_size'}})

• By setting the ONNX version to 10 i.e (opset_version=10)

Here we specify the ONNX version to export the model to. The version that supports the specified slicing operation is 10. So we add another parameter to the exporter that specifies this version. i.e.

opset_version=10

So our exporter should be:

with tempfile.TemporaryDirectory() as tmpdir:
    model_path = os.path.join(tmpdir,"customAdder.onnx") 
    torch.onnx.export(model,x_tensor,model_path,
                      export_params=True,opset_version=10,
                      do_constant_folding=True,
                      input_names = ['input'],
                      output_names = ['output'],
                      dynamic_axes={'input' : {0 : 'batch_size'},
                                    'output' : {0 : 'batch_size'}})

By setting opset_version >= 10, slice operation will be supported by torch.onnx