Real-ESRGAN

1
import argparse
2
import torch
3
import torch.onnx
4
from basicsr.archs.rrdbnet_arch import RRDBNet
5

6

7
def main(args):
8
    # An instance of the model
9
    model = RRDBNet(num_in_ch=3, num_out_ch=3, num_feat=64, num_block=23, num_grow_ch=32, scale=4)
10
    if args.params:
11
        keyname = 'params'
12
    else:
13
        keyname = 'params_ema'
14
    model.load_state_dict(torch.load(args.input)[keyname])
15
    # set the train mode to false since we will only run the forward pass.
16
    model.train(False)
17
    model.cpu().eval()
18

19
    # An example input
20
    x = torch.rand(1, 3, 64, 64)
21
    # Export the model
22
    with torch.no_grad():
23
        torch_out = torch.onnx._export(model, x, args.output, opset_version=11, export_params=True)
24
    print(torch_out.shape)
25

26

27
if __name__ == '__main__':
28
    """Convert pytorch model to onnx models"""
29
    parser = argparse.ArgumentParser()
30
    parser.add_argument(
31
        '--input', type=str, default='experiments/pretrained_models/RealESRGAN_x4plus.pth', help='Input model path')
32
    parser.add_argument('--output', type=str, default='realesrgan-x4.onnx', help='Output onnx path')
33
    parser.add_argument('--params', action='store_false', help='Use params instead of params_ema')
34
    args = parser.parse_args()
35

36
    main(args)
37