adapt Real-ESRGAN-anime model

This commit is contained in:
Xintao
2021-08-31 22:28:30 +08:00
parent c1669c4b0a
commit 6ff747174d
3 changed files with 16 additions and 4 deletions

View File

@@ -3,7 +3,7 @@ import torch.onnx
from basicsr.archs.rrdbnet_arch import RRDBNet
# An instance of your model
model = RRDBNet(num_in_ch=3, num_out_ch=3, num_feat=64, num_block=23, num_grow_ch=32)
model = RRDBNet(num_in_ch=3, num_out_ch=3, num_feat=64, num_block=23, num_grow_ch=32, scale=4)
model.load_state_dict(torch.load('experiments/pretrained_models/RealESRGAN_x4plus.pth')['params_ema'])
# set the train mode to false since we will only run the forward pass.
model.train(False)