PyTorch ONNX export

[albertz]

This issue is to track any aspects and issues on PyTorch (#1120) ONNX export.

• Working script for conversion (export_to_onnx.py)
• Fix issue with convolution
• Rename script to torch_export_to_onnx.py
• Check model_outputs.
• Use model_outputs dims when not specified in mark_as_output (esp in case of PT).
• The input names for sizes should be better (currently they are like "classes:size1" but should be like "data:size1"). (Fixed via Torch ONNX export, proper dynamic dims handling #1362.)
• The batch dim (data:size0), or any scalar dyn_size_ext, should not be an input, as it is redundant. (Fixed via Torch ONNX export, proper dynamic dims handling #1362.)
• model_outputs: Ignore dim tag equality. Check static dims, dyn dims information
• Out seq lens are actually relevant, and should not be filled with dummy values? Clarify. (PyTorch ONNX export #1333 (comment)) (Fixed via Torch ONNX export, proper dynamic dims handling #1362.)
• Working demo-rf + test case
• Test case: Also perform ONNX inference (Fixed via Torch ONNX export, proper dynamic dims handling #1362.)
• RF Conformer works
• Some real-world pure PT model works
• Running the resulting ONNX model in RASR
• Writing a sisyphus job to run ONNX export in i6_core. (Fixed via Add OnnxExportJob i6_core#429.)

Other things which are less important for now:

• Avoiding TracerWarnings for some sanity checking code: I think it's not really possible in an easy way.

[albertz]

The initial version is done already (export_to_onnx.py) and works for pure PT code, so I think the main work is done. We can already close the issue. But let's anyway add further aspects on this here.

[albertz]

One aspect is that convolution with padding="same" is not supported by ONNX. (See also here: #1120 (comment)) However, we can simply work around that by doing the padding manually. The code is already there for the RF, and just needs to be used. We can use torch.onnx.is_in_export_mode() to test whether we should use it.

[Icemole]

As a form of summary: the tool is mainly done, but there are some operators that aren't yet supported, one of which is the convolutional layer padding when specified as a string ("same" or "valid"). This prevents the demos from being run with the tool, which is one of my top priorities.

I found that it's not actually padding="same" which is not implemented, it's actually type(padding) is str what fails: a string can't be provided as padding, since I tried padding="valid" (what we had already discussed) and it failed with the same message, even when padding=0 worked fine.

I've been investigating this issue and I'm currently trying to fix it. There's some people online having the same issues and they also provide implementations of custom padding, so it shouldn't be very hard to fix. I'll try to push a version with my changes before the weekend.

[albertz]

padding="valid" is just the same as padding=0, or not?

[Icemole]

Yes it is, but it seems that internally any padding given as a string triggers some parameter called aten::_convolution_mode, which isn't implemented in the ONNX conversion.

Anyway, padding="valid" is the easy case to fix, since then we only have to set padding = 0, as you said. padding="same" is a little more tricky in my view: it involves not only calculating the padding for each dimension (which is the easy part) but also ordering it properly in the result.

Also, padding is given as a tuple of maximum 3 dimensions ([[depth,] height,] width) depending on the dimensions of the convolutional layer (3D, 2D or 1D), so I have to investigate a bit the internals in order to check if there's redundant padding inserted that we don't desire: for instance, with a kernel size of 2x2 and a stride of 1, I think we'd only need to insert an additional row and column, but maybe this tuple considers symmetric padding and inserts one row at the start and another at the end, and the same for the column, making the output size not match the input size.

[albertz]

Anyway, padding="valid" is the easy case to fix, since then we only have to set padding = 0, as you said.

Yes, my question was just to confirm this. So we can just always do this then, no matter if ONNX or so. Just:

if padding == "valid":
    padding = 0

padding="same" is a little more tricky in my view: it involves not only calculating the padding for each dimension (which is the easy part) but also ordering it properly in the result.

We already have the code for adding the padding manual. See the code. It basically does:

if padding == "same" and <manual_padding_required>:
    ...  # add padding to input
    padding = "valid"

Nothing else needs to be done, or not?

[albertz]

Btw, next step after the current demo works: Test our RF Conformer implementation.

[albertz]

Actually, I just wonder now: Instead of manually adding the padding to the input, could we maybe just set padding = (filter_size - 1) // 2? But probably not really in all cases. E.g. only works with odd filter_size. No striding. Etc. But it might be worth to still have this case, as it might be faster than the manual padding.

[Icemole]

I was able to insert torch.onnx.is_in_export_mode() into the already developed code which creates manual striding. However, this only works for the RF demo: the PT demo as it is can't be executed, since it also uses padding="same" but we have no handler for pure PT code. Should we just let these details to the user, or also map "same" to a valid padding size?

[Icemole]

could we maybe just set padding = (filter_size - 1) // 2?

Internally, the forward() implementation of convolutional layers is as follows:

        if self.padding_mode != 'zeros':
            return F.conv1d(F.pad(input, self._reversed_padding_repeated_twice, mode=self.padding_mode),
                            weight, bias, self.stride,
                            _single(0), self.dilation, self.groups)

where self._reversed_padding_repeated_twice([i, j]) is [i, i, j, j]. This variable is actually calculated at initialization time. I think if we also do that and introduce an attribute like this one in our convolutional layers, we can save quite a bit of operations as we wouldn't have to calculate the same paddings for every forward, right?

Or maybe we could cache it somehow, or add an additional layer attribute self.padding_sizes at runtime?

[albertz]

I was able to insert torch.onnx.is_in_export_mode() into the already developed code which creates manual striding. However, this only works for the RF demo: the PT demo as it is can't be executed, since it also uses padding="same" but we have no handler for pure PT code. Should we just let these details to the user, or also map "same" to a valid padding size?

Change padding="same" to padding=2 in the Torch demo. That should fix this issue.

[albertz]

Internally, the forward() implementation of convolutional layers is as follows ...

This is only for the case padding_mode != 'zeros', which we don't support at all currently, right?