Here's what I did to export the pth
to onnx
(if only for my own future reference):
- Open the Colab notebook and click Runtime > Run All.
- Open up and
main_test_swinir.py
in the Colab editor and placing the following line afteroutput = model(img_lq)
:
torch.onnx.export(model, img_lq, "003_realSR_BSRGAN_DFO_s64w8_SwinIR-M_x4_GAN-dynamic.onnx", export_params=True, opset_version=12, do_constant_folding=True, verbose=True, input_names = ['input'], output_names = ['output'], dynamic_axes={'input' : {2 : 'h', 3 : 'w'}, 'output' : {2 : 'h', 3 : 'w'}})
- Run this:
!python SwinIR/main_test_swinir.py --task real_sr --model_path experiments/pretrained_models/003_realSR_BSRGAN_DFO_s64w8_SwinIR-M_x4_GAN.pth --folder_lq BSRGAN/testsets/RealSRSet --scale 4
And the ONNX file that you see in this repo will be generated.