swin-ir-onnx / README.md
rocca's picture
Update README.md
63b5441
metadata
license: apache-2.0

Here's what I did to export the pth to onnx (if only for my own future reference):

  1. Open the Colab notebook and click Runtime > Run All.
  2. Open up and main_test_swinir.py in the Colab editor and placing the following line after output = 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'}})
  1. 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.