pth2onnx,How should I use “torch.onnx.export()”

[diplomatist]

pth2onnx,How should I use “torch.onnx.export()”

[diplomatist]

The current idea is to save the parameters of the forward process after converting the denoise_fn module of the network, which is unet, to onnx.

In the process of exporting the onnx model, it appears
runtime error: couldn't export python mehod. but no specific part was given as to what went wrong.

I tested each module of the network and found that one of the modules used the checkpoint function, which I removed and successfully exported the model in onnx format.

However, the exported ONNX model actually increases the inference time during inference.
There may be several reasons:
① The checkpoint function in unet has been removed by me.
② The input of onnx only supports numpy format, which means that the input parameters of the model can only be in the CPU environment.
Perhaps there will be better improvements only after converting onnx to TensorRT. However, recent DDIM, PLMS, and dpm solver's inverse Markov chain sampling methods have greatly accelerated the sampling process of the diffusion model. The inference time gain brought by the new sampling method has far exceeded the gain brought by TensorRT deployment.