ESRGAN/net_interp.py

22 lines
567 B
Python
Raw Normal View History

2018-09-05 18:21:30 +03:00
import sys
import torch
from collections import OrderedDict
alpha = float(sys.argv[1])
net_PSNR_path = './models/RRDB_PSNR_x4.pth'
net_ESRGAN_path = './models/RRDB_ESRGAN_x4.pth'
2018-09-05 18:28:53 +03:00
net_interp_path = './models/interp_{:02d}.pth'.format(int(alpha*10))
2018-09-05 18:21:30 +03:00
net_PSNR = torch.load(net_PSNR_path)
net_ESRGAN = torch.load(net_ESRGAN_path)
net_interp = OrderedDict()
2018-09-05 18:28:53 +03:00
print('Interpolating with alpha = ', alpha)
2018-09-05 18:21:30 +03:00
for k, v_PSNR in net_PSNR.items():
v_ESRGAN = net_ESRGAN[k]
2018-09-05 18:28:53 +03:00
net_interp[k] = (1 - alpha) * v_PSNR + alpha * v_ESRGAN
2018-09-05 18:21:30 +03:00
2018-09-05 18:28:53 +03:00
torch.save(net_interp, net_interp_path)