def train(args, imgs):
# Load model
net = skip(num_input_channels=32,
num_output_channels=3,
num_channels_down=[64, 64, 64, 128, 128],
num_channels_up=[64, 64, 64, 128, 128],
num_channels_skip=[4, 4, 4, 4, 4],
upsample_mode='bilinear',
need_sigmoid=True,
need_bias=True,
pad='reflection',
act_fun='LeakyReLU')
net = net.float()
net = net.cuda() if torch.cuda.is_available() else net
# Compute the number of parameters
s = sum([np.prod(list(p.size())) for p in net.parameters()])
print('Number of parameters: ', s)
criterion = nn.MSELoss()
# define input
net_input = get_noise(32, 'noise', (np.shape(
imgs['HR_np'])[1], np.shape(imgs['HR_np'])[2])).float().detach()
img_LR_var = np_to_var(imgs['LR_np']).float()
downsampler = Downsampler(n_planes=3,
factor=args.factor,
kernel_type='lanczos2',
phase=0.5,
preserve_size=True).float()
downsampler = downsampler.cuda() if torch.cuda.is_available(
) else downsampler
net_input_saved = net_input.data.clone()
noise = net_input.data.clone()
# Define closure
reg_noise_std = 0.0
tv_weight = 0.0
net_input = net_input.cuda() if torch.cuda.is_available() else net_input
img_LR_var = img_LR_var.cuda() if torch.cuda.is_available() else img_LR_var
net_input_saved = net_input_saved.cuda() if torch.cuda.is_available(
) else net_input_saved
def closure():
global iteration, psnr_HR
if reg_noise_std > 0:
net_input.data = net_input_saved + (
torch.cuda.FloatTensor(noise.size()).normal_() * reg_noise_std)
out_HR = net(net_input)
out_LR = downsampler(out_HR)
total_loss = criterion(out_LR, img_LR_var)
if tv_weight > 0:
total_loss += tv_weight * tv_loss(out_HR)
total_loss.backward()
# Log
psnr_LR = compare_psnr(imgs['LR_np'], out_LR.data.cpu().numpy()[0])
psnr_HR = compare_psnr(imgs['HR_np'], out_HR.data.cpu().numpy()[0])
print('Iteration %05d Loss %3f PSNR_LR %.3f PSNR_HR %.3f' %
(iteration, total_loss.data[0], psnr_LR, psnr_HR),
'\r',
end='')
if iteration % 500 == 0 or iteration == args.epoch - 1:
out_HR_np = var_to_np(out_HR)
# plot_image_grid([imgs['HR_np'], imgs['bicubic_np'], np.clip(out_HR_np, 0, 1)], factor=13, nrow=3)
io.imsave(
'sr_' + str(iteration) + '.png',
np.transpose(np.clip(out_HR_np, 0, 1), [1, 2, 0])[:, :, :3])
iteration += 1
return total_loss
# Optimize (Adam with first 100 epoch and LBFGS with rest)
p = get_params('net', net, net_input)
optimizer = torch.optim.Adam(net.parameters(), lr=0.001)
for j in range(100):
optimizer.zero_grad()
closure()
optimizer.step()
print('Starting optimization with LBFGS')
if args.epoch > 100:
optimizer = torch.optim.LBFGS(net.parameters(),
max_iter=args.epoch - 100,
lr=0.01,
tolerance_grad=-1,
tolerance_change=-1)
def closure2():
optimizer.zero_grad()
return closure()
optimizer.step(closure2)
# Show final result
out_HR_np = np.clip(var_to_np(net(net_input)), 0, 1)
result_deep_prior = put_in_center(out_HR_np, imgs['HR_np'].shape[1:])
plot_image_grid([imgs['HR_np'], imgs['bicubic_np'], out_HR_np],
factor=4,
nrow=1)
print('\nFinal PSNR: ', psnr_HR)
'skip',
pad,
skip_n33d=128,
skip_n33u=128,
skip_n11=4,
num_scales=5,
upsample_mode='bilinear').type(dtype)
# Losses
mse = torch.nn.MSELoss().type(dtype)
img_LR_var = np_to_torch(imgs['LR_np']).type(dtype)
downsampler = Downsampler(n_planes=3,
factor=factor,
kernel_type=KERNEL_TYPE,
phase=0.5,
preserve_size=True).type(dtype)
# # Define closure and optimize
# In[ ]:
def closure():
global i, net_input
if reg_noise_std > 0:
net_input = net_input_saved + (noise.normal_() * reg_noise_std)
out_HR = net(net_input)
def __init__(self, config):
self.rank, self.world_size = 0, 1
if config['dist']:
self.rank = dist.get_rank()
self.world_size = dist.get_world_size()
self.config = config
self.mode = config['dgp_mode']
self.update_G = config['update_G']
self.update_embed = config['update_embed']
self.iterations = config['iterations']
self.ftr_num = config['ftr_num']
self.ft_num = config['ft_num']
self.lr_ratio = config['lr_ratio']
self.G_lrs = config['G_lrs']
self.z_lrs = config['z_lrs']
self.use_in = config['use_in']
self.select_num = config['select_num']
self.factor = 2 if self.mode == 'hybrid' else 4 # Downsample factor
# create model
self.G = models.Generator(**config).cuda()
self.D = models.Discriminator(
**config).cuda() if config['ftr_type'] == 'Discriminator' else None
self.G.optim = torch.optim.Adam(
[{
'params': self.G.get_params(i, self.update_embed)
} for i in range(len(self.G.blocks) + 1)],
lr=config['G_lr'],
betas=(config['G_B1'], config['G_B2']),
weight_decay=0,
eps=1e-8)
# load weights
if config['random_G']:
self.random_G()
else:
utils.load_weights(self.G if not (config['use_ema']) else None,
self.D,
config['weights_root'],
name_suffix=config['load_weights'],
G_ema=self.G if config['use_ema'] else None,
strict=False)
self.G.eval()
if self.D is not None:
self.D.eval()
self.G_weight = deepcopy(self.G.state_dict())
# prepare latent variable and optimizer
self._prepare_latent()
# prepare learning rate scheduler
self.G_scheduler = utils.LRScheduler(self.G.optim, config['warm_up'])
self.z_scheduler = utils.LRScheduler(self.z_optim, config['warm_up'])
# loss functions
self.mse = torch.nn.MSELoss()
if config['ftr_type'] == 'Discriminator':
self.ftr_net = self.D
self.criterion = utils.DiscriminatorLoss(
ftr_num=config['ftr_num'][0])
else:
vgg = torchvision.models.vgg16(pretrained=True).cuda().eval()
self.ftr_net = models.subsequence(vgg.features, last_layer='20')
self.criterion = utils.PerceptLoss()
# Downsampler for producing low-resolution image
self.downsampler = Downsampler(n_planes=3,
factor=self.factor,
kernel_type='lanczos2',
phase=0.5,
preserve_size=True).type(
torch.cuda.FloatTensor)
from models.model_sr import Model
net = Model()
net = net.type(dtype)
net_input = get_noise(
args.input_depth, args.noise_method,
(imgs['HR_pil'].size[1], imgs['HR_pil'].size[0])).type(dtype).detach()
mse = torch.nn.MSELoss().type(dtype)
img_LR_var = np_to_torch(imgs['LR_np']).type(dtype)
downsampler = Downsampler(n_planes=3,
factor=args.factor,
kernel_type='lanczos2',
phase=0.5,
preserve_size=True).type(dtype)
psnr_gt_best = 0
i = 0
PSNR_list = []
_t = {'im_detect': Timer(), 'misc': Timer()}
def closure():
global i, net_input, psnr_gt_best, PSNR_list
_t['im_detect'].tic()
def get_h(n_ch, blur_type, use_fourier, dtype):
assert blur_type in ['uniform_blur', 'gauss_blur'], "blur_type can be or 'uniform' or 'gauss'"
if not use_fourier:
return Downsampler(n_ch, 1, blur_type, preserve_size=True).type(dtype)
return lambda im: torch_blur(im, blur_type, dtype)