Create Real-ESRGAN_x2plus__convert_pth_to_onnx.py

Real-ESRGAN_x2plus__convert_pth_to_onnx.py

@@ -0,0 +1,32 @@
+import torch
+import torch.onnx as onnx
+from basicsr.archs.rrdbnet_arch import RRDBNet
+
+# Load the PyTorch model
+device = torch.device('cpu')
+model = RRDBNet(num_in_ch=3, num_out_ch=3, num_feat=64, num_block=23, num_grow_ch=32, scale=2)
+
+# Load the state dictionary
+state_dict = torch.load('Real-ESRGAN_x2plus.pth', map_location=device)
+
+# Load the state dictionary
+model.load_state_dict(state_dict['params_ema'])
+model.train(False)
+
+# Set the model to evaluation mode
+model.eval()
+
+# Define the input shape
+input_shape = (1, 3, 64, 64)  # batch_size, channels, height, width
+
+# Create a dummy input tensor
+dummy_input = torch.randn(input_shape)
+
+# Convert the model to ONNX
+onnx.export(model, 
+            dummy_input, 
+            'Real-ESRGAN_x2plus.onnx', 
+            opset_version=11, 
+            input_names=['input'], 
+            output_names=['output'], 
+            dynamic_axes={'input': {0: 'batch_size'}, 'output': {0: 'batch_size'}})