def test(self):
# networks
self.num_recursions = 16
self.model = Net(num_channels=self.num_channels,
base_filter=256,
num_recursions=self.num_recursions)
if self.gpu_mode:
self.model.cuda()
# load model
self.load_model()
# load dataset
test_data_loader = self.load_dataset(dataset='test')
# Test
print('Test is started.')
img_num = 0
self.model.eval()
for input, target in test_data_loader:
# input data (bicubic interpolated image)
if self.gpu_mode:
y_ = Variable(
utils.img_interp(input, self.scale_factor).cuda())
else:
y_ = Variable(utils.img_interp(input, self.scale_factor))
# prediction
_, recon_imgs = self.model(y_)
for i, recon_img in enumerate(recon_imgs):
img_num += 1
recon_img = recon_imgs[i].cpu().data
gt_img = target[i]
lr_img = input[i]
bc_img = utils.img_interp(input[i], self.scale_factor)
# calculate psnrs
bc_psnr = utils.PSNR(bc_img, gt_img)
recon_psnr = utils.PSNR(recon_img, gt_img)
# save result images
result_imgs = [gt_img, lr_img, bc_img, recon_img]
psnrs = [None, None, bc_psnr, recon_psnr]
utils.plot_test_result(result_imgs,
psnrs,
img_num,
save_dir=self.save_dir)
print("Saving %d test result images..." % img_num)
def test(self):
# networks
self.model = Net(num_channels=self.num_channels, base_filter=64, scale_factor=self.scale_factor)
if self.gpu_mode:
self.model.cuda()
# load model
self.load_model()
# load dataset
test_data_loader = self.load_dataset(dataset='test')
# Test
print('Test is started.')
img_num = 0
self.model.eval()
for input, target in test_data_loader:
# input data (low resolution image)
if self.gpu_mode:
y_ = Variable(input.cuda())
else:
y_ = Variable(input)
# prediction
recon_imgs = self.model(y_)
for i in range(self.test_batch_size):
img_num += 1
recon_img = recon_imgs[i].cpu().data
gt_img = utils.shave(target[i], border_size=8 * self.scale_factor)
lr_img = utils.shave(input[i], border_size=8)
bc_img = utils.shave(utils.img_interp(input[i], self.scale_factor), border_size=8 * self.scale_factor)
# calculate psnrs
bc_psnr = utils.PSNR(bc_img, gt_img)
recon_psnr = utils.PSNR(recon_img, gt_img)
# save result images
result_imgs = [gt_img, lr_img, bc_img, recon_img]
psnrs = [None, None, bc_psnr, recon_psnr]
utils.plot_test_result(result_imgs, psnrs, img_num, save_dir=self.save_dir)
print("Saving %d test result images..." % img_num)
(epoch + 1, params.num_epochs, i + 1, len(train_data_loader),
D_loss.data[0], G_loss.data[0]))
D_avg_loss = torch.mean(torch.FloatTensor(D_losses))
G_avg_loss = torch.mean(torch.FloatTensor(G_losses))
# avg loss values for plot
D_avg_losses.append(D_avg_loss)
G_avg_losses.append(G_avg_loss)
# Show result for test image
gen_image = G(Variable(test_input)) #.cuda()))
gen_image = gen_image.cpu().data
utils.plot_test_result(test_input,
test_target,
gen_image,
epoch,
save=True,
save_dir=save_dir)
np.savetxt("D_losses.csv", D_losses, delimiter=",")
np.savetxt("G_losses.csv", G_losses, delimiter=",")
# Plot average losses
utils.plot_loss(D_losses,
G_losses,
params.num_epochs,
save=True,
save_dir=save_dir)
# Make gif
utils.make_gif(params.dataset, params.num_epochs, save_dir=save_dir)
def test(self):
# networks
self.G = Generator(num_channels=self.num_channels,
base_filter=64,
num_residuals=16)
if self.gpu_mode:
self.G.cuda()
# load model
self.load_model()
# load dataset
for test_dataset in self.test_dataset:
test_data_loader = self.load_dataset(dataset=test_dataset,
is_train=False)
# Test
print('Test is started.')
img_num = 0
total_img_num = len(test_data_loader)
self.G.eval()
for lr, hr, bc in test_data_loader:
# input data (low resolution image)
if self.num_channels == 1:
y_ = Variable(utils.norm(lr[:, 0].unsqueeze(1), vgg=True))
else:
y_ = Variable(utils.norm(lr, vgg=True))
if self.gpu_mode:
y_ = y_.cuda()
# prediction
recon_imgs = self.G(y_)
for i, recon_img in enumerate(recon_imgs):
img_num += 1
sr_img = utils.denorm(recon_img.cpu().data, vgg=True)
# save result image
save_dir = os.path.join(self.save_dir, 'test_result',
test_dataset)
utils.save_img(sr_img, img_num, save_dir=save_dir)
# calculate psnrs
if self.num_channels == 1:
gt_img = hr[i][0].unsqueeze(0)
lr_img = lr[i][0].unsqueeze(0)
bc_img = bc[i][0].unsqueeze(0)
else:
gt_img = hr[i]
lr_img = lr[i]
bc_img = bc[i]
bc_psnr = utils.PSNR(bc_img, gt_img)
recon_psnr = utils.PSNR(sr_img, gt_img)
# plot result images
result_imgs = [gt_img, lr_img, bc_img, sr_img]
psnrs = [None, None, bc_psnr, recon_psnr]
utils.plot_test_result(result_imgs,
psnrs,
img_num,
save_dir=save_dir)
print('Test DB: %s, Saving result images...[%d/%d]' %
(test_dataset, img_num, total_img_num))
print('Test is finishied.')
def train(self):
# load dataset
train_data_loader = self.load_dataset(dataset=self.train_dataset,
is_train=True)
test_data_loader = self.load_dataset(dataset=self.test_dataset[0],
is_train=False)
# networks
self.G = Generator(num_channels=self.num_channels,
base_filter=64,
num_residuals=16)
self.D = Discriminator(num_channels=self.num_channels,
base_filter=64,
image_size=self.crop_size)
# weigh initialization
self.G.weight_init()
self.D.weight_init()
# For the content loss
self.feature_extractor = FeatureExtractor(
models.vgg19(pretrained=True))
# optimizer
self.G_optimizer = optim.Adam(self.G.parameters(),
lr=self.lr,
betas=(0.9, 0.999))
# self.D_optimizer = optim.Adam(self.D.parameters(), lr=self.lr, betas=(0.9, 0.999))
self.D_optimizer = optim.SGD(self.D.parameters(),
lr=self.lr / 100,
momentum=0.9,
nesterov=True)
# loss function
if self.gpu_mode:
self.G.cuda()
self.D.cuda()
self.feature_extractor.cuda()
self.MSE_loss = nn.MSELoss().cuda()
self.BCE_loss = nn.BCELoss().cuda()
else:
self.MSE_loss = nn.MSELoss()
self.BCE_loss = nn.BCELoss()
print('---------- Networks architecture -------------')
utils.print_network(self.G)
utils.print_network(self.D)
print('----------------------------------------------')
# set the logger
G_log_dir = os.path.join(self.save_dir, 'G_logs')
if not os.path.exists(G_log_dir):
os.mkdir(G_log_dir)
G_logger = Logger(G_log_dir)
D_log_dir = os.path.join(self.save_dir, 'D_logs')
if not os.path.exists(D_log_dir):
os.mkdir(D_log_dir)
D_logger = Logger(D_log_dir)
################# Pre-train generator #################
self.epoch_pretrain = 50
# Load pre-trained parameters of generator
if not self.load_model(is_pretrain=True):
# Pre-training generator for 50 epochs
print('Pre-training is started.')
self.G.train()
for epoch in range(self.epoch_pretrain):
for iter, (lr, hr, _) in enumerate(train_data_loader):
# input data (low resolution image)
if self.num_channels == 1:
x_ = Variable(
utils.norm(hr[:, 0].unsqueeze(1), vgg=True))
y_ = Variable(
utils.norm(lr[:, 0].unsqueeze(1), vgg=True))
else:
x_ = Variable(utils.norm(hr, vgg=True))
y_ = Variable(utils.norm(lr, vgg=True))
if self.gpu_mode:
x_ = x_.cuda()
y_ = y_.cuda()
# Train generator
self.G_optimizer.zero_grad()
recon_image = self.G(y_)
# Content losses
content_loss = self.MSE_loss(recon_image, x_)
# Back propagation
G_loss_pretrain = content_loss
G_loss_pretrain.backward()
self.G_optimizer.step()
# log
print("Epoch: [%2d] [%4d/%4d] G_loss_pretrain: %.8f" %
((epoch + 1), (iter + 1), len(train_data_loader),
G_loss_pretrain.item()))
print('Pre-training is finished.')
# Save pre-trained parameters of generator
self.save_model(is_pretrain=True)
################# Adversarial train #################
print('Training is started.')
# Avg. losses
G_avg_loss = []
D_avg_loss = []
step = 0
# test image
test_lr, test_hr, test_bc = test_data_loader.dataset.__getitem__(2)
test_lr = test_lr.unsqueeze(0)
test_hr = test_hr.unsqueeze(0)
test_bc = test_bc.unsqueeze(0)
self.G.train()
self.D.train()
for epoch in range(self.num_epochs):
# learning rate is decayed by a factor of 10 every 20 epoch
if (epoch + 1) % 20 == 0:
for param_group in self.G_optimizer.param_groups:
param_group["lr"] /= 10.0
print("Learning rate decay for G: lr={}".format(
self.G_optimizer.param_groups[0]["lr"]))
for param_group in self.D_optimizer.param_groups:
param_group["lr"] /= 10.0
print("Learning rate decay for D: lr={}".format(
self.D_optimizer.param_groups[0]["lr"]))
G_epoch_loss = 0
D_epoch_loss = 0
for iter, (lr, hr, _) in enumerate(train_data_loader):
# input data (low resolution image)
mini_batch = lr.size()[0]
if self.num_channels == 1:
x_ = Variable(utils.norm(hr[:, 0].unsqueeze(1), vgg=True))
y_ = Variable(utils.norm(lr[:, 0].unsqueeze(1), vgg=True))
else:
x_ = Variable(utils.norm(hr, vgg=True))
y_ = Variable(utils.norm(lr, vgg=True))
if self.gpu_mode:
x_ = x_.cuda()
y_ = y_.cuda()
# labels
real_label = Variable(torch.ones(mini_batch).cuda())
fake_label = Variable(torch.zeros(mini_batch).cuda())
else:
# labels
real_label = Variable(torch.ones(mini_batch))
fake_label = Variable(torch.zeros(mini_batch))
# Reset gradient
self.D_optimizer.zero_grad()
# Train discriminator with real data
D_real_decision = self.D(x_)
D_real_loss = self.BCE_loss(D_real_decision[:, 0], real_label)
# Train discriminator with fake data
recon_image = self.G(y_)
D_fake_decision = self.D(recon_image)
D_fake_loss = self.BCE_loss(D_fake_decision[:, 0], fake_label)
D_loss = D_real_loss + D_fake_loss
# Back propagation
D_loss.backward()
self.D_optimizer.step()
# Reset gradient
self.G_optimizer.zero_grad()
# Train generator
recon_image = self.G(y_)
D_fake_decision = self.D(recon_image)
# Adversarial loss
GAN_loss = self.BCE_loss(D_fake_decision[:, 0], real_label)
# Content losses
mse_loss = self.MSE_loss(recon_image, x_)
x_VGG = Variable(utils.norm(hr, vgg=True).cuda())
recon_VGG = Variable(
utils.norm(recon_image.data, vgg=True).cuda())
real_feature = self.feature_extractor(x_VGG)
fake_feature = self.feature_extractor(recon_VGG)
vgg_loss = self.MSE_loss(fake_feature, real_feature.detach())
# Back propagation
G_loss = mse_loss + 6e-3 * vgg_loss + 1e-3 * GAN_loss
G_loss.backward()
self.G_optimizer.step()
# log
G_epoch_loss += G_loss.item()
D_epoch_loss += D_loss.item()
print("Epoch: [%2d] [%4d/%4d] G_loss: %.8f, D_loss: %.8f" %
((epoch + 1), (iter + 1), len(train_data_loader),
G_loss.item(), D_loss.item()))
# tensorboard logging
#G_logger.scalar_summary('losses', G_loss.item(), step + 1)
#D_logger.scalar_summary('losses', D_loss.item(), step + 1)
step += 1
# avg. loss per epoch
G_avg_loss.append(G_epoch_loss / len(train_data_loader))
D_avg_loss.append(D_epoch_loss / len(train_data_loader))
# prediction
if self.num_channels == 1:
y_ = Variable(utils.norm(test_lr[:, 0].unsqueeze(1), vgg=True))
else:
y_ = Variable(utils.norm(test_lr, vgg=True))
if self.gpu_mode:
y_ = y_.cuda()
recon_img = self.G(y_)
sr_img = utils.denorm(recon_img[0].cpu().data, vgg=True)
# save result image
save_dir = os.path.join(self.save_dir, 'train_result')
utils.save_img(sr_img,
epoch + 1,
save_dir=save_dir,
is_training=True)
print('Result image at epoch %d is saved.' % (epoch + 1))
# Save trained parameters of model
if (epoch + 1) % self.save_epochs == 0:
self.save_model(epoch + 1)
# calculate psnrs
if self.num_channels == 1:
gt_img = test_hr[0][0].unsqueeze(0)
lr_img = test_lr[0][0].unsqueeze(0)
bc_img = test_bc[0][0].unsqueeze(0)
else:
gt_img = test_hr[0]
lr_img = test_lr[0]
bc_img = test_bc[0]
bc_psnr = utils.PSNR(bc_img, gt_img)
recon_psnr = utils.PSNR(sr_img, gt_img)
# plot result images
result_imgs = [gt_img, lr_img, bc_img, sr_img]
psnrs = [None, None, bc_psnr, recon_psnr]
utils.plot_test_result(result_imgs,
psnrs,
self.num_epochs,
save_dir=save_dir,
is_training=True)
print('Training result image is saved.')
# Plot avg. loss
utils.plot_loss([G_avg_loss, D_avg_loss],
self.num_epochs,
save_dir=self.save_dir)
print("Training is finished.")
# Save final trained parameters of model
self.save_model(epoch=None)
for i, real_A in enumerate(test_data_loader_A):
print (len(real_A))
# input image data
real_A=real_A[0]
print (len(real_A))
print (real_A.shape)
real_A = Variable(gpuAvailable(real_A, params.cuda))
# A -> B -> A
fake_B = G_A(real_A)
recon_A = G_B(fake_B)
# Show result for test data
utils.plot_test_result(real_A, fake_B, recon_A, i, save=True, save_dir=save_dir + 'AtoB/')
print('%d images are generated.' % (i + 1))
for i, real_B in enumerate(test_data_loader_B):
# input image data
real_B = real_B[0]
real_B = Variable(gpuAvailable(real_B, params.cuda))
# B -> A -> B
fake_A = G_B(real_B)
recon_B = G_A(fake_A)
# Show result for test data
utils.plot_test_result(real_B, fake_A, recon_B, i, save=True, save_dir=save_dir + 'BtoA/')
def train(self):
# networks
self.model = Net(num_channels=self.num_channels, base_filter=64)
# weigh initialization
self.model.weight_init(mean=0.0, std=0.001)
# optimizer
self.optimizer = optim.SGD(self.model.parameters(), lr=self.lr)
# loss function
if self.gpu_mode:
self.model.cuda()
self.MSE_loss = nn.MSELoss().cuda()
else:
self.MSE_loss = nn.MSELoss()
print('---------- Networks architecture -------------')
utils.print_network(self.model)
print('----------------------------------------------')
# load dataset
train_data_loader = self.load_dataset(dataset='train')
test_data_loader = self.load_dataset(dataset='test')
# set the logger
log_dir = os.path.join(self.save_dir, 'logs')
if not os.path.exists(log_dir):
os.mkdir(log_dir)
logger = Logger(log_dir)
################# Train #################
print('Training is started.')
avg_loss = []
step = 0
# test image
test_input, test_target = test_data_loader.dataset.__getitem__(2)
test_input = test_input.unsqueeze(0)
test_target = test_target.unsqueeze(0)
self.model.train()
for epoch in range(self.num_epochs):
epoch_loss = 0
for iter, (input, target) in enumerate(train_data_loader):
# input data (bicubic interpolated image)
if self.gpu_mode:
# exclude border pixels from loss computation
x_ = Variable(utils.shave(target, border_size=8).cuda())
y_ = Variable(utils.img_interp(input, self.scale_factor).cuda())
else:
x_ = Variable(utils.shave(target, border_size=8))
y_ = Variable(utils.img_interp(input, self.scale_factor))
# update network
self.optimizer.zero_grad()
recon_image = self.model(y_)
loss = self.MSE_loss(recon_image, x_)
loss.backward()
self.optimizer.step()
# log
epoch_loss += loss.data[0]
print("Epoch: [%2d] [%4d/%4d] loss: %.8f" % ((epoch + 1), (iter + 1), len(train_data_loader), loss.data[0]))
# tensorboard logging
logger.scalar_summary('loss', loss.data[0], step + 1)
step += 1
# avg. loss per epoch
avg_loss.append(epoch_loss / len(train_data_loader))
# prediction
recon_imgs = self.model(Variable(utils.img_interp(test_input, self.scale_factor).cuda()))
recon_img = recon_imgs[0].cpu().data
gt_img = utils.shave(test_target[0], border_size=8)
lr_img = test_input[0]
bc_img = utils.shave(utils.img_interp(test_input[0], self.scale_factor), border_size=8)
# calculate psnrs
bc_psnr = utils.PSNR(bc_img, gt_img)
recon_psnr = utils.PSNR(recon_img, gt_img)
# save result images
result_imgs = [gt_img, lr_img, bc_img, recon_img]
psnrs = [None, None, bc_psnr, recon_psnr]
utils.plot_test_result(result_imgs, psnrs, epoch + 1, save_dir=self.save_dir, is_training=True)
print("Saving training result images at epoch %d" % (epoch + 1))
# Save trained parameters of model
if (epoch + 1) % self.save_epochs == 0:
self.save_model(epoch + 1)
# Plot avg. loss
utils.plot_loss([avg_loss], self.num_epochs, save_dir=self.save_dir)
print("Training is finished.")
# Save final trained parameters of model
self.save_model(epoch=None)
def train(self):
# networks
self.model = Net(num_channels=self.num_channels,
base_filter=64,
num_residuals=18)
# weigh initialization
self.model.weight_init()
# optimizer
self.momentum = 0.9
self.weight_decay = 0.0001
self.clip = 0.4
self.optimizer = optim.SGD(self.model.parameters(),
lr=self.lr,
momentum=self.momentum,
weight_decay=self.weight_decay)
# loss function
if self.gpu_mode:
self.model.cuda()
self.MSE_loss = nn.MSELoss().cuda()
else:
self.MSE_loss = nn.MSELoss()
print('---------- Networks architecture -------------')
utils.print_network(self.model)
print('----------------------------------------------')
# load dataset
train_data_loader = self.load_dataset(dataset='train')
test_data_loader = self.load_dataset(dataset='test')
# set the logger
log_dir = os.path.join(self.save_dir, 'logs')
if not os.path.exists(log_dir):
os.mkdir(log_dir)
logger = Logger(log_dir)
################# Train #################
print('Training is started.')
avg_loss = []
step = 0
# test image
test_input, test_target, test_bicubic = test_data_loader.dataset.__getitem__(
2)
test_input = test_input.unsqueeze(0)
test_target = test_target.unsqueeze(0)
self.model.train()
for epoch in range(self.num_epochs):
# learning rate is decayed by a factor of 10 every 20 epochs
if (epoch + 1) % 20 == 0:
for param_group in self.optimizer.param_groups:
param_group["lr"] /= 10.0
print("Learning rate decay: lr={}".format(
self.optimizer.param_groups[0]["lr"]))
epoch_loss = 0
for iter, (input, target, bi) in enumerate(train_data_loader):
# input data (bicubic interpolated image)
if self.gpu_mode:
x_ = Variable(target.cuda())
y_ = Variable(
utils.img_interp(input, self.scale_factor).cuda())
else:
x_ = Variable(target)
y_ = Variable(utils.img_interp(input, self.scale_factor))
# update network
self.optimizer.zero_grad()
recon_image = self.model(y_)
loss = self.MSE_loss(recon_image, x_)
loss.backward()
# gradient clipping
nn.utils.clip_grad_norm(self.model.parameters(), self.clip)
self.optimizer.step()
# log
epoch_loss += loss.item()
print("Epoch: [%2d] [%4d/%4d] loss: %.8f" %
((epoch + 1),
(iter + 1), len(train_data_loader), loss.item()))
# tensorboard logging
#logger.scalar_summary('loss', loss.item(), step + 1)
step += 1
# avg. loss per epoch
avg_loss.append(epoch_loss / len(train_data_loader))
# prediction
recon_imgs = self.model(
Variable(
utils.img_interp(test_input, self.scale_factor).cuda()))
recon_img = recon_imgs[0].cpu().data
gt_img = test_target[0]
lr_img = test_input[0]
bc_img = utils.img_interp(test_input[0], self.scale_factor)
# calculate psnrs
bc_psnr = utils.PSNR(bc_img, gt_img)
recon_psnr = utils.PSNR(recon_img, gt_img)
# save result images
result_imgs = [gt_img, lr_img, bc_img, recon_img]
psnrs = [None, None, bc_psnr, recon_psnr]
utils.plot_test_result(result_imgs,
psnrs,
epoch + 1,
save_dir=self.save_dir,
is_training=True)
print("Saving training result images at epoch %d" % (epoch + 1))
# Save trained parameters of model
if (epoch + 1) % self.save_epochs == 0:
self.save_model(epoch + 1)
# Plot avg. loss
utils.plot_loss([avg_loss], self.num_epochs, save_dir=self.save_dir)
print("Training is finished.")
# Save final trained parameters of model
self.save_model(epoch=None)
direction=params.direction,
transform=test_transform)
test_data_loader = torch.utils.data.DataLoader(dataset=test_data,
batch_size=params.batch_size,
shuffle=False)
# Load model
G = Generator(3, params.ngf, 3)
G.cuda()
G.load_state_dict(torch.load(model_dir + 'generator_param.pkl'))
# Test
for i, (input, target) in enumerate(test_data_loader):
# input & target image data
x_ = Variable(input.cuda())
y_ = Variable(target.cuda())
gen_image = G(x_)
gen_image = gen_image.cpu().data
# Show result for test data
utils.plot_test_result(input,
target,
gen_image,
i,
training=False,
save=True,
save_dir=save_dir)
print('%d images are generated.' % (i + 1))
def train(self):
# networks
self.num_recursions = 16
self.model = Net(num_channels=self.num_channels,
base_filter=256,
num_recursions=self.num_recursions)
# weigh initialization
self.model.weight_init()
# optimizer
self.momentum = 0.9
self.weight_decay = 0.0001
self.loss_alpha = 1.0
self.loss_alpha_zero_epoch = 25
self.loss_alpha_decay = self.loss_alpha / self.loss_alpha_zero_epoch
self.loss_beta = 0.001
# learnable parameters
param_groups = list(self.model.parameters())
param_groups = [{'params': param_groups}]
param_groups += [{'params': [self.model.w]}]
self.optimizer = optim.Adam(param_groups, lr=self.lr)
# loss function
if self.gpu_mode:
self.model.cuda()
self.MSE_loss = nn.MSELoss().cuda()
else:
self.MSE_loss = nn.MSELoss()
print('---------- Networks architecture -------------')
utils.print_network(self.model)
print('----------------------------------------------')
# load dataset
train_data_loader = self.load_dataset(dataset='train')
test_data_loader = self.load_dataset(dataset='test')
# set the logger
log_dir = os.path.join(self.save_dir, 'logs')
if not os.path.exists(log_dir):
os.mkdir(log_dir)
logger = Logger(log_dir)
################# Train #################
print('Training is started.')
avg_loss = []
step = 0
# test image
test_input, test_target = test_data_loader.dataset.__getitem__(2)
test_input = test_input.unsqueeze(0)
test_target = test_target.unsqueeze(0)
self.model.train()
for epoch in range(self.num_epochs):
# learning rate is decayed by a factor of 10 every 20 epochs
if (epoch + 1) % 20 == 0:
for param_group in self.optimizer.param_groups:
param_group["lr"] /= 10.0
print("Learning rate decay: lr={}".format(
self.optimizer.param_groups[0]["lr"]))
# loss_alpha decayed to zero after 25 epochs
self.loss_alpha = max(0.0, self.loss_alpha - self.loss_alpha_decay)
epoch_loss = 0
for iter, (input, target) in enumerate(train_data_loader):
# input data (bicubic interpolated image)
if self.gpu_mode:
y = Variable(target.cuda())
x = Variable(
utils.img_interp(input, self.scale_factor).cuda())
else:
y = Variable(target)
x = Variable(utils.img_interp(input, self.scale_factor))
# update network
self.optimizer.zero_grad()
y_d_, y_ = self.model(x)
# loss1
loss1 = 0
for d in range(self.num_recursions):
loss1 += (self.MSE_loss(y_d_[d], y) / self.num_recursions)
# loss2
loss2 = self.MSE_loss(y_, y)
# regularization
reg_term = 0
for theta in self.model.parameters():
reg_term += torch.mean(torch.sum(theta**2))
# total loss
loss = self.loss_alpha * loss1 + (
1 - self.loss_alpha) * loss2 + self.loss_beta * reg_term
loss.backward()
self.optimizer.step()
# log
epoch_loss += loss.data[0]
print("Epoch: [%2d] [%4d/%4d] loss: %.8f" %
((epoch + 1),
(iter + 1), len(train_data_loader), loss.data[0]))
# tensorboard logging
logger.scalar_summary('loss', loss.data[0], step + 1)
step += 1
# avg. loss per epoch
avg_loss.append(epoch_loss / len(train_data_loader))
# prediction
_, recon_imgs = self.model(
Variable(
utils.img_interp(test_input, self.scale_factor).cuda()))
recon_img = recon_imgs[0].cpu().data
gt_img = test_target[0]
lr_img = test_input[0]
bc_img = utils.img_interp(test_input[0], self.scale_factor)
# calculate psnrs
bc_psnr = utils.PSNR(bc_img, gt_img)
recon_psnr = utils.PSNR(recon_img, gt_img)
# save result images
result_imgs = [gt_img, lr_img, bc_img, recon_img]
psnrs = [None, None, bc_psnr, recon_psnr]
utils.plot_test_result(result_imgs,
psnrs,
epoch + 1,
save_dir=self.save_dir,
is_training=True)
print("Saving training result images at epoch %d" % (epoch + 1))
# Save trained parameters of model
if (epoch + 1) % self.save_epochs == 0:
self.save_model(epoch + 1)
# Plot avg. loss
utils.plot_loss([avg_loss], self.num_epochs, save_dir=self.save_dir)
print("Training is finished.")
# Save final trained parameters of model
self.save_model(epoch=None)
def train(self):
# networks
self.model = Net(3, 64, 3, 1)
LossNet.creat_loss_Net(self)
if self.gpu_mode:
self.model.cuda()
if self.pre_epochs == 0:
# weigh initialization
self.model.weight_init()
else:
# (self, is training or not, is E_model(True) or D_model(False))
utils.load_model(self, True, True)
# self.model.load_state_dict(torch.load(self.E_pretrain))
# optimizer
self.optimizer = optim.Adam(self.model.parameters(), lr=self.lr)
self.valid = None
self.fake = None
if 'A' in self.model_loss:
self.discriminator = Discriminator().cuda(
) if self.gpu_mode else Discriminator()
Tensor = torch.cuda.FloatTensor if self.gpu_mode else torch.Tensor
self.optimizer_D = torch.optim.Adam(
self.discriminator.parameters(), lr=0.0002, betas=(0.5, 0.999))
if self.pre_epochs != 0:
utils.load_model(self, True, False)
if self.loss_F == "BCEWithLogitsLoss":
self.criterion_GAN = nn.BCEWithLogitsLoss(
size_average=False).cuda(
) if self.gpu_mode else nn.BCEWithLogitsLoss(
size_average=False)
# elif self.loss_F == "Cross"
# print('---------- Networks architecture -------------')
# utils.print_network(self.model)
# print('----------------------------------------------')
# load dataset
train_data_loader = self.load_dataset(dataset=self.train_dataset,
is_train=True)
test_data_loader = self.load_dataset(dataset=self.test_dataset[0],
is_train=False)
# set the logger
log_dir = os.path.join(self.save_dir, 'logs')
if not os.path.exists(log_dir):
os.makedirs(log_dir)
logger = Logger(log_dir)
################# Train #################
print('Training is started.')
avg_loss = []
step = 0
avg_loss_D = []
# test image
test_lr, test_hr, test_bc, name = test_data_loader.dataset.__getitem__(
2)
test_lr = test_lr.unsqueeze(0)
test_hr = test_hr.unsqueeze(0)
test_bc = test_bc.unsqueeze(0)
self.model.train()
for epoch in range(self.num_epochs):
# learning rate is decayed by a factor of 2 every 40 epochs
if (epoch + 1) % 40 == 0:
for param_group in self.optimizer.param_groups:
param_group['lr'] /= 2.0
print('Learning rate decay: lr={}'.format(
self.optimizer.param_groups[0]['lr']))
epoch_loss = 0
epoch_loss_D = 0
for iter, (lr, hr, bc_lr) in enumerate(train_data_loader):
# input data (low resolution image)
x_ = Variable(hr).cuda() if self.gpu_mode else Variable(hr)
y_ = Variable(lr).cuda() if self.gpu_mode else Variable(hr)
bc_y_ = Variable(bc_lr).cuda() if self.gpu_mode else Variable(
hr)
recon_image = self.model(y_)
recon_image = recon_image + bc_y_
loss_G = 0
loss_D = 0
style_loss = 0
loss_T = []
if 'A' in self.model_loss:
# ---------------------
# Train Discriminator
# ---------------------
self.valid = Variable(Tensor(
np.ones(x_.size()[0]).reshape((x_.size()[0], 1))),
requires_grad=False)
self.fake = Variable(Tensor(
np.zeros(x_.size()[0]).reshape((x_.size()[0], 1))),
requires_grad=False)
if self.gpu_mode:
self.valid.cuda()
self.fake.cuda()
self.optimizer_D.zero_grad()
# Loss of real and fake images
if self.loss_F == "BCEWithLogitsLoss":
loss_real = self.criterion_GAN(self.discriminator(x_),
self.valid)
loss_fake = self.criterion_GAN(
self.discriminator(recon_image.detach()),
self.fake)
elif self.loss_F == "MSE":
loss_real = self.mse_loss(self.discriminator(x_),
self.valid)
loss_fake = self.mse_loss(
self.discriminator(recon_image.detach()),
self.fake)
# Total loss
loss_D = (loss_real + loss_fake) / 2
loss_D.backward()
self.optimizer_D.step()
epoch_loss_D += loss_D.data[0]
if iter % self.D_period == 0:
# -----------------
# Train Generator
# -----------------
self.optimizer.zero_grad()
loss_a, loss_output_m2, loss_output_m5, style_score, loss_G, loss_T = Loss.loss_op(
self, self, recon_image, x_)
loss = (2*0.1*loss_output_m2) + \
(2*0.01*loss_output_m5) + \
loss_a*2 + style_score*1e-6
loss.backward()
self.optimizer.step()
# log
epoch_loss += loss.data[0]
utils.print_loss(self, epoch, len(train_data_loader),
loss, style_score, loss_output_m2,
loss_output_m5, iter, loss_a, loss_D,
loss_G, loss_T)
# tensorboard logging
logger.scalar_summary('loss', loss.data[0], step + 1)
step += 1
del x_, y_, bc_y_
if self.gpu_mode:
torch.cuda.empty_cache()
else:
# update network
self.optimizer.zero_grad()
loss_a, loss_output_m2, loss_output_m5, style_score, loss_G, loss_T = Loss.loss_op(
self, self, recon_image, x_)
loss = (2*0.1*loss_output_m2) + \
(2*0.01*loss_output_m5) + \
loss_a*1e-3 + style_score*1e-6
loss.backward()
self.optimizer.step()
# log
epoch_loss += loss.data[0]
#epoch_loss += loss
utils.print_loss(self, epoch, len(train_data_loader), loss,
style_score, loss_output_m2,
loss_output_m5, iter, loss_a, loss_D,
loss_G, loss_T)
# tensorboard logging
logger.scalar_summary('loss', loss.data[0], step + 1)
step += 1
del x_, y_, bc_y_
if self.gpu_mode:
torch.cuda.empty_cache()
# avg. loss per epoch
avg_loss.append(epoch_loss / len(train_data_loader))
# prediction
y_ = Variable(test_lr).cuda() if self.gpu_mode else Variable(
test_lr)
bc = Variable(test_bc).cuda() if self.gpu_mode else Variable(
test_bc)
recon_img = self.model(y_)
recon_img = recon_img + bc
sr_img = recon_img[0].cpu().data
# save result image
save_dir = os.path.join(self.save_dir, 'train_result')
utils.save_img(sr_img,
epoch + 1,
save_dir=save_dir,
is_training=True)
print('Result image at epoch %d is saved.' % (epoch + 1))
#
utils.plot_loss([avg_loss], epoch, save_dir=save_dir)
if 'A' in self.model_loss:
avg_loss_D.append(epoch_loss_D / len(train_data_loader))
utils.plot_loss([avg_loss_D], epoch, save_dir=save_dir + '/D')
del y_, bc
if self.gpu_mode:
torch.cuda.empty_cache()
# Save trained parameters of model
if (epoch + 1) % self.save_epochs == 0:
utils.save_model(self, epoch + 1)
test_v2.save_TrainingTest(self, epoch + 1)
# calculate psnrs
if self.num_channels == 1:
gt_img = test_hr[0][0].unsqueeze(0)
lr_img = test_lr[0][0].unsqueeze(0)
bc_img = test_bc[0][0].unsqueeze(0)
else:
gt_img = test_hr[0]
lr_img = test_lr[0]
bc_img = test_bc[0]
bc_psnr = utils.PSNR(bc_img, gt_img)
recon_psnr = utils.PSNR(sr_img, gt_img)
# plot result images
result_imgs = [gt_img, lr_img, bc_img, sr_img]
psnrs = [None, None, bc_psnr, recon_psnr]
utils.plot_test_result(result_imgs,
psnrs,
self.num_epochs,
save_dir=save_dir,
is_training=True)
print('Training result image is saved.')
# Plot avg. loss
utils.plot_loss([avg_loss], self.num_epochs, save_dir=save_dir)
print('Training is finished.')
# Save final trained parameters of model
utils.save_model(self, epoch=None)
def test(self):
# networks
self.model = Net(3, 64, 3, 1)
# load model
utils.load_model(self, False, False)
# self.model.load_state_dict(torch.load('epoch0_70.pkl'))
if self.gpu_mode:
self.model.cuda()
# load dataset
for test_dataset in self.test_dataset:
test_data_loader = self.load_dataset(dataset=test_dataset,
is_train=False)
# Test
print('Test is started.')
img_num = 0
total_img_num = len(test_data_loader)
self.model.eval()
for lr, hr, bc, _ in test_data_loader:
# input data (low resolution image)
x_ = Variable(hr).cuda() if self.gpu_mode else Variable(hr)
y_ = Variable(lr).cuda() if self.gpu_mode else Variable(hr)
bc_ = Variable(bc).cuda() if self.gpu_mode else Variable(hr)
# prediction
recon_imgs = self.model(y_)
recon_imgs += bc_
for i, recon_img in enumerate(recon_imgs):
img_num += 1
sr_img = recon_img.cpu().data
# save result image
# save_dir = os.path.join(
# self.save_dir, 'test_result_texture', test_dataset)
save_dir = os.path.join(self.save_dir, 'test_result',
test_dataset)
utils.save_img(sr_img, img_num, save_dir=save_dir)
# utils.save_img(x_, img_num, save_dir=os.path.join(
# self.save_dir, 'test_HR', test_dataset))
# calculate psnrs
if self.num_channels == 1:
gt_img = hr[i][0].unsqueeze(0)
lr_img = lr[i][0].unsqueeze(0)
bc_img = bc[i][0].unsqueeze(0)
else:
gt_img = hr[i]
lr_img = lr[i]
bc_img = bc[i]
bc_psnr = utils.PSNR(bc_img, gt_img)
recon_psnr = utils.PSNR(sr_img, gt_img)
# plot result images
result_imgs = [gt_img, lr_img, bc_img, sr_img]
psnrs = [None, None, bc_psnr, recon_psnr]
utils.plot_test_result(result_imgs,
psnrs,
img_num,
save_dir=save_dir)
print('Test DB: %s, Saving result images...[%d/%d]' %
(test_dataset, img_num, total_img_num))
print('Test is finishied.')
def train(self):
# networks
self.model = Net(num_channels=self.num_channels,
base_filter=64,
num_residuals=16)
# weigh initialization
self.model.weight_init()
# optimizer
self.optimizer = optim.Adam(self.model.parameters(),
lr=self.lr,
betas=(0.9, 0.999),
eps=1e-8)
# loss function
if self.gpu_mode:
self.model.cuda()
self.L1_loss = nn.L1Loss().cuda()
else:
self.L1_loss = nn.L1Loss()
print('---------- Networks architecture -------------')
utils.print_network(self.model)
print('----------------------------------------------')
# load dataset
train_data_loader = self.load_dataset(dataset=self.train_dataset,
is_train=True)
test_data_loader = self.load_dataset(dataset=self.test_dataset[0],
is_train=False)
# set the logger
log_dir = os.path.join(self.save_dir, 'logs')
if not os.path.exists(log_dir):
os.makedirs(log_dir)
logger = Logger(log_dir)
################# Train #################
print('Training is started.')
avg_loss = []
step = 0
# test image
test_lr, test_hr, test_bc = test_data_loader.dataset.__getitem__(2)
test_lr = test_lr.unsqueeze(0)
test_hr = test_hr.unsqueeze(0)
test_bc = test_bc.unsqueeze(0)
self.model.train()
for epoch in range(self.num_epochs):
# learning rate is decayed by a factor of 2 every 40 epochs
if (epoch + 1) % 40 == 0:
for param_group in self.optimizer.param_groups:
param_group['lr'] /= 2.0
print('Learning rate decay: lr={}'.format(
self.optimizer.param_groups[0]['lr']))
epoch_loss = 0
for iter, (lr, hr, _) in enumerate(train_data_loader):
# input data (low resolution image)
if self.num_channels == 1:
x_ = Variable(hr[:, 0].unsqueeze(1))
y_ = Variable(lr[:, 0].unsqueeze(1))
else:
x_ = Variable(hr)
y_ = Variable(lr)
if self.gpu_mode:
x_ = x_.cuda()
y_ = y_.cuda()
# update network
self.optimizer.zero_grad()
recon_image = self.model(y_)
loss = self.L1_loss(recon_image, x_)
loss.backward()
self.optimizer.step()
# log
epoch_loss += loss.data[0]
print('Epoch: [%2d] [%4d/%4d] loss: %.8f' %
((epoch + 1),
(iter + 1), len(train_data_loader), loss.data[0]))
# tensorboard logging
logger.scalar_summary('loss', loss.data[0], step + 1)
step += 1
# avg. loss per epoch
avg_loss.append(epoch_loss / len(train_data_loader))
# prediction
if self.num_channels == 1:
y_ = Variable(test_lr[:, 0].unsqueeze(1))
else:
y_ = Variable(test_lr)
if self.gpu_mode:
y_ = y_.cuda()
recon_img = self.model(y_)
sr_img = recon_img[0].cpu().data
# save result image
save_dir = os.path.join(self.save_dir, 'train_result')
utils.save_img(sr_img,
epoch + 1,
save_dir=save_dir,
is_training=True)
print('Result image at epoch %d is saved.' % (epoch + 1))
# Save trained parameters of model
if (epoch + 1) % self.save_epochs == 0:
self.save_model(epoch + 1)
# calculate psnrs
if self.num_channels == 1:
gt_img = test_hr[0][0].unsqueeze(0)
lr_img = test_lr[0][0].unsqueeze(0)
bc_img = test_bc[0][0].unsqueeze(0)
else:
gt_img = test_hr[0]
lr_img = test_lr[0]
bc_img = test_bc[0]
bc_psnr = utils.PSNR(bc_img, gt_img)
recon_psnr = utils.PSNR(sr_img, gt_img)
# plot result images
result_imgs = [gt_img, lr_img, bc_img, sr_img]
psnrs = [None, None, bc_psnr, recon_psnr]
utils.plot_test_result(result_imgs,
psnrs,
self.num_epochs,
save_dir=save_dir,
is_training=True)
print('Training result image is saved.')
# Plot avg. loss
utils.plot_loss([avg_loss], self.num_epochs, save_dir=save_dir)
print('Training is finished.')
# Save final trained parameters of model
self.save_model(epoch=None)
