FFT support for tensorrt version of pytorch/tensorflow models

Hi All,
My model has some FFT and IFFT operations in between of my deep learning model. I would like to convert the model to Tensorrt for high throughput inference. I have seen the supported layers https://docs.nvidia.com/deeplearning/sdk/tensorrt-support-matrix/index.html and it doesn’t mention anything about FFT operations. I am new to tensorrt (Cuda aswell)and I would like to know that if I need to write FFT and IFFT plugins ? or can the problem be solved by just using CuFFT on output from one network layer and feeding the output from CuFFT to the next layer ?

Any help is appreciated. Thank You!


Do you have a way of converting your full model to ONNX format? If so, you can try to convert to TensorRT and see what falls out of that. TensorRT 7.0 was just released with expanded ONNX support.

Can you also elaborate more on your deep learning model (framework, language, etc.), and how you’re connecting the FFT operations in between?

Thank you for your reply. My model is in Pytorch 1.2.0 and I have some FFT and IFFT layers in my model which we use to convert our Image to Frequency domain and back. In the Frequency domain we perform some operations to preserve certain frequencies and then use IFFT to convert back to image domain and continue with regular deep learning. It is a RNN kind of an architecture where FTT and IFFT happen at the end of RNN blocks and this continues for N number of time. I have already tried converting my pytorch model to ONNX but ONNX doesn’t support FFT layers yet. So I thought my safest bet would be directly writing the network in tensorrt 6.0 and now I am stuck with the FFT and IFFT part of the network and I don’t know what would be the best solution here.

If I write a plugin which takes in data and does FFT and IFFT using CuFFT lib, do you think it can work ? I don’t have much experience with Cuda or Tensorrt so I am not sure if this will work. Any suggestions is appreciated. Thank you!


As a long term solution, it might be a good idea to reach out to the PyTorch team at https://github.com/pytorch/pytorch/issues regarding future ONNX support for your FFT/IFFT ops since they’re built into PyTorch. They might have some useful insight at the least - feel free to link that issue here so I can follow it’s status if any.

For a current/short-term solution, I would think you’re right to say that you could implement them as custom plugins into your TensorRT network. Hopefully some of the plugin samples in the TensorRT release as well as in the TensorRT OSS components might be able to help with the plugin side of things: https://github.com/NVIDIA/TensorRT/tree/master/samples/opensource/samplePlugin. And maybe this repo could help with the CuFFT implementation: https://github.com/NVIDIA/cuda-samples/tree/master/Samples/simpleCUFFT