add onnx export

[HaiFengZeng]

b = 2
c = 64
L = 125
step =3
stride = 2
inputs = torch.randn((b, c, 1, L))
lib_module = unfoldNd.UnfoldNd(
    (1, step), dilation=(1, 1), padding=0, stride=stride
)
lib_outputs = lib_module(inputs)
a2 = lib_outputs.view(b, c, step, -1).transpose(2, 3)
unfold = x.unfold(-1, 3, 2)

torch.onnx.export(lib_module, inputs, 'model.onnx',
                  input_names=['x'],
                  output_names=['y'],
                  opset_version=14,
                  dynamic_axes={'x': {3: 'frame'}}
                  )

great work!The trick using convolutions seems good and I get the same result as unfold = x.unfold(-1, 3, 2), But when I try to export to onnx, I got this error: RuntimeError: Unsupported: ONNX export of convolution for kernel of unknown shape. is there any idea to export onnx with the dynamic convolution size ? thanks!

[f-dangel]

Hi, thanks for the interesting request!

I believe the 'problem' for exporting with ONNX is that the one-hot kernels are created during the forward pass. However, their size is known beforehand. So putting together a ONNX-exportable version would be possible: Just extract the kernel generation from the forward(...) into the __init__ method. Probably one would also have to change the interface to pass the required information for constructing the one-hot kernels in the constructor.

Therefore, I am not sure if such a module can be part of the main lib without duplicating all the code. Any suggestion how to adapt the current interface of UnfoldNd would be welcome.

Could you make the example executable? I can then try to put together a solution.

Best, Felix

[HaiFengZeng]

Thank you! That's good advice, and I can export to onnx now, just modify : The kernel size and the input channel are used in make weight as convolution kernel, so we make it a initialized parameters

    def __init__(self, kernel_size, c_in, dilation=1, padding=0, stride=1):
        super().__init__()

        self._kernel_size = kernel_size
        self._dilation = dilation
        self._padding = padding
        self._stride = stride
        self.c_in = c_in

    def forward(self, input):
        return unfoldNd(
            input,
            self._kernel_size,
            c_in=self.c_in,
            dilation=self._dilation,
            padding=self._padding,
            stride=self._stride,
        )

and batch_size, in_channels = input.shape[0], c_in in unfoldNd function

[f-dangel]

That's great!

For now, I won't integrate this solution into the main library unless there is a huge demand for ONNX export. However, it would be great to have a self-contained code snippet for other users that may run into this problem.

Could you prepare and post an MWE (basically the modified UnfoldNd class and the ONNX export) in this thread? That would be extremely helpful :)

Best, Felix

[f-dangel]

Closing as stale.