RuntimeError: Exporting the operator col2im to ONNX opset version 12 is not supported

[didadida-r]

hi, i convert the stft model to onnx version , but a error is raised.
'''
RuntimeError: Exporting the operator col2im to ONNX opset version 12 is not supported. Please feel free to request support or submit a pull request on PyTorch GitHub
'''

[KinWaiCheuk]

At first glance, it seems this error is related to PyTorch version, rather than nnAudio. Which PyTorch version are you using?
Could you please share me the code for converting the stft model to onnx?
I will try to locate the exact problem, thanks!

[KinWaiCheuk]

I ran a quick test on my computer with the following code, and I can obtain the onnx file successfully.

from nnAudio.Spectrogram import STFT
import torch

dummy_input = torch.randn(4,44100)
stft_layer = STFT()
torch.onnx.export(stft_layer,
                  dummy_input,
                  'stft_layer.onnx',
                  verbose=True,
                  input_names=['waveforms'],
                  output_names=['stft_output'])

I am using pytorch==1.9.0+cu111 and nnAudio==0.2.5. Could you double check if you are able to do ONNX using the same pytorch and nnAudio as me?

[didadida-r]

hi, i test it again and this problem might locate in the iSTFT module, here is the test code.

    from nnAudio.Spectrogram import STFT, iSTFT
    import torch
    
    class DummpyModel(nn.Module):
        def __init__(self):
            super(DummpyModel, self).__init__()
            self.stft_layer = STFT()
            self.istft_layer = iSTFT()
        
        def forward(self, wav):
            out = self.stft_layer(wav, output_format="Complex")
            out = self.istft_layer(out, length=wav.shape[-1], onesided=True)
            # out_wav = self.decoder(spec, wav.shape)
            return out

    dummy_input = torch.randn(4,44100)
    model = DummpyModel()
    torch.onnx.export(model,
                    dummy_input,
                    'stft_layer.onnx',
                    verbose=True,
                    opset_version=12,
                    input_names=['waveforms'],
                    output_names=['stft_output'])

[KinWaiCheuk]

You are right, when trying to export iSTFT to onnx, it will result in an error message:

RuntimeError: Exporting the operator col2im to ONNX opset version 12 is not supported. Please feel free to request support or submit a pull request on PyTorch GitHub.

This is because the fold function that iSTFT uses is not yet supported by onnx.

nnAudio/Installation/nnAudio/utils.py

Lines 19 to 29 in ac9cd13

	def torch_window_sumsquare(w, n_frames, stride, n_fft, power=2):
	w_stacks = w.unsqueeze(-1).repeat((1,n_frames)).unsqueeze(0)
	# Window length + stride*(frames-1)
	output_len = w_stacks.shape[1] + stride*(w_stacks.shape[2]-1)
	return fold(w_stacks**power, (1,output_len), kernel_size=(1,n_fft), stride=stride)
	def overlap_add(X, stride):
	n_fft = X.shape[1]
	output_len = n_fft + stride*(X.shape[2]-1)
	return fold(X, (1,output_len), kernel_size=(1,n_fft), stride=stride).flatten(1)

There is a open issue under the PyTorch repo, I think people are working on it already.
pytorch/pytorch#41423

Unfortunately, there is not much I can do at the moment. One workaround is to implement the two functions, torch_window_sumsquare and overlap_add without using fold. I do have an old version of nnAudio that did not use fold for inverse before commit 14db7a6e94e3f38ccaded7c779a85b99fdaefff1. But I remember the inverse was quite slow without using fold, and the result was slightly incorrect.

[veronicamorelli]

Operator col2im is supported for opset version 18 but not the previous ones, is they a way to overcome the error?
"torch.onnx.errors.UnsupportedOperatorError: Exporting the operator 'aten::col2im' to ONNX opset version 17 is not supported. Support for this operator was added in version 18, try exporting with this version."