def LSGAN_trainer(opt):
# ----------------------------------------
# Initialize training parameters
# ----------------------------------------
# cudnn benchmark accelerates the network
if opt.cudnn_benchmark == True:
cudnn.benchmark = True
else:
cudnn.benchmark = False
# Build networks
generator = utils.create_generator(opt)
discriminator = utils.create_discriminator(opt)
perceptualnet = utils.create_perceptualnet()
# To device
if opt.multi_gpu == True:
generator = nn.DataParallel(generator)
discriminator = nn.DataParallel(discriminator)
perceptualnet = nn.DataParallel(perceptualnet)
generator = generator.cuda()
discriminator = discriminator.cuda()
perceptualnet = perceptualnet.cuda()
else:
generator = generator.cuda()
discriminator = discriminator.cuda()
perceptualnet = perceptualnet.cuda()
# Loss functions
L1Loss = nn.L1Loss()
MSELoss = nn.MSELoss()
#FeatureMatchingLoss = FML1Loss(opt.fm_param)
# Optimizers
optimizer_g = torch.optim.Adam(generator.parameters(),
lr=opt.lr_g,
betas=(opt.b1, opt.b2),
weight_decay=opt.weight_decay)
optimizer_d = torch.optim.Adam(generator.parameters(),
lr=opt.lr_d,
betas=(opt.b1, opt.b2),
weight_decay=opt.weight_decay)
# Learning rate decrease
def adjust_learning_rate(lr_in, optimizer, epoch, opt):
"""Set the learning rate to the initial LR decayed by "lr_decrease_factor" every "lr_decrease_epoch" epochs"""
lr = lr_in * (opt.lr_decrease_factor**(epoch // opt.lr_decrease_epoch))
for param_group in optimizer.param_groups:
param_group['lr'] = lr
# Save the model if pre_train == True
def save_model(net, epoch, opt):
"""Save the model at "checkpoint_interval" and its multiple"""
if opt.multi_gpu == True:
if epoch % opt.checkpoint_interval == 0:
torch.save(
net.module, 'deepfillNet_epoch%d_batchsize%d.pth' %
(epoch, opt.batch_size))
print('The trained model is successfully saved at epoch %d' %
(epoch))
else:
if epoch % opt.checkpoint_interval == 0:
torch.save(
net, 'deepfillNet_epoch%d_batchsize%d.pth' %
(epoch, opt.batch_size))
print('The trained model is successfully saved at epoch %d' %
(epoch))
# ----------------------------------------
# Initialize training dataset
# ----------------------------------------
# Define the dataset
trainset = dataset.InpaintDataset(opt)
print('The overall number of images equals to %d' % len(trainset))
# Define the dataloader
dataloader = DataLoader(trainset,
batch_size=opt.batch_size,
shuffle=True,
num_workers=opt.num_workers,
pin_memory=True)
# ----------------------------------------
# Training and Testing
# ----------------------------------------
# Initialize start time
prev_time = time.time()
# Tensor type
Tensor = torch.cuda.FloatTensor
# Training loop
for epoch in range(opt.epochs):
for batch_idx, (img, mask) in enumerate(dataloader):
# Load mask (shape: [B, 1, H, W]), masked_img (shape: [B, 3, H, W]), img (shape: [B, 3, H, W]) and put it to cuda
img = img.cuda()
mask = mask.cuda()
# LSGAN vectors
valid = Tensor(np.ones((img.shape[0], 1, 8, 8)))
fake = Tensor(np.zeros((img.shape[0], 1, 8, 8)))
### Train Discriminator
optimizer_d.zero_grad()
# Generator output
first_out, second_out = generator(img, mask)
# forward propagation
first_out_wholeimg = img * (
1 - mask) + first_out * mask # in range [-1, 1]
second_out_wholeimg = img * (
1 - mask) + second_out * mask # in range [-1, 1]
# Fake samples
fake_scalar = discriminator(second_out_wholeimg.detach(), mask)
# True samples
true_scalar = discriminator(img, mask)
# Overall Loss and optimize
loss_fake = MSELoss(fake_scalar, fake)
loss_true = MSELoss(true_scalar, valid)
# Overall Loss and optimize
loss_D = 0.5 * (loss_fake + loss_true)
loss_D.backward()
optimizer_d.step()
### Train Generator
optimizer_g.zero_grad()
# Mask L1 Loss
first_MaskL1Loss = L1Loss(first_out_wholeimg, img)
second_MaskL1Loss = L1Loss(second_out_wholeimg, img)
# GAN Loss
fake_scalar = discriminator(second_out_wholeimg, mask)
GAN_Loss = MSELoss(fake_scalar, valid)
# Get the deep semantic feature maps, and compute Perceptual Loss
img = (img + 1) / 2 # in range [0, 1]
img = utils.normalize_ImageNet_stats(img) # in range of ImageNet
img_featuremaps = perceptualnet(img) # feature maps
second_out_wholeimg = (second_out_wholeimg +
1) / 2 # in range [0, 1]
second_out_wholeimg = utils.normalize_ImageNet_stats(
second_out_wholeimg)
second_out_wholeimg_featuremaps = perceptualnet(
second_out_wholeimg)
second_PerceptualLoss = L1Loss(second_out_wholeimg_featuremaps,
img_featuremaps)
# Compute losses
loss = first_MaskL1Loss + second_MaskL1Loss + opt.perceptual_param * second_PerceptualLoss + opt.gan_param * GAN_Loss
loss.backward()
optimizer_g.step()
# Determine approximate time left
batches_done = epoch * len(dataloader) + batch_idx
batches_left = opt.epochs * len(dataloader) - batches_done
time_left = datetime.timedelta(seconds=batches_left *
(time.time() - prev_time))
prev_time = time.time()
# Print log
print(
"\r[Epoch %d/%d] [Batch %d/%d] [first Mask L1 Loss: %.5f] [second Mask L1 Loss: %.5f]"
% ((epoch + 1), opt.epochs, batch_idx, len(dataloader),
first_MaskL1Loss.item(), second_MaskL1Loss.item()))
print(
"\r[D Loss: %.5f] [G Loss: %.5f] [Perceptual Loss: %.5f] time_left: %s"
% (loss_D.item(), GAN_Loss.item(),
second_PerceptualLoss.item(), time_left))
# Learning rate decrease
adjust_learning_rate(opt.lr_g, optimizer_g, (epoch + 1), opt)
adjust_learning_rate(opt.lr_d, optimizer_d, (epoch + 1), opt)
# Save the model
save_model(generator, (epoch + 1), opt)
### warp I1 and O1
warp_i1 = flow_warping(frame_i1, flow_i21)
warp_o1 = flow_warping(frame_o1, flow_i21)
### compute non-occlusion mask: exp(-alpha * || F_i2 - Warp(F_i1) ||^2 )
noc_mask2 = torch.exp( -opts.alpha * torch.sum(frame_i2 - warp_i1, dim=1).pow(2) ).unsqueeze(1)
ST_loss += opts.w_ST * criterion(frame_o2 * noc_mask2, warp_o1 * noc_mask2)
### perceptual loss
if opts.w_VGG > 0:
### normalize
frame_o2_n = utils.normalize_ImageNet_stats(frame_o2)
frame_p2_n = utils.normalize_ImageNet_stats(frame_p2)
### extract VGG features
features_p2 = VGG(frame_p2_n, opts.VGGLayers[-1])
features_o2 = VGG(frame_o2_n, opts.VGGLayers[-1])
VGG_loss_all = []
for l in opts.VGGLayers:
VGG_loss_all.append( criterion(features_o2[l], features_p2[l]) )
VGG_loss += opts.w_VGG * sum(VGG_loss_all)
## end of forward
def Trainer_GAN(opt):
# ----------------------------------------
# Initialization
# ----------------------------------------
# cudnn benchmark
cudnn.benchmark = opt.cudnn_benchmark
# configurations
save_folder = os.path.join(opt.save_path, opt.task_name)
sample_folder = os.path.join(opt.sample_path, opt.task_name)
if not os.path.exists(save_folder):
os.makedirs(save_folder)
if not os.path.exists(sample_folder):
os.makedirs(sample_folder)
# Initialize networks
generator = utils.create_generator(opt)
discriminator = utils.create_discriminator(opt)
perceptualnet = utils.create_perceptualnet()
# To device
if opt.multi_gpu:
generator = nn.DataParallel(generator)
generator = generator.cuda()
discriminator = nn.DataParallel(discriminator)
discriminator = discriminator.cuda()
perceptualnet = nn.DataParallel(perceptualnet)
perceptualnet = perceptualnet.cuda()
else:
generator = generator.cuda()
discriminator = discriminator.cuda()
perceptualnet = perceptualnet.cuda()
# ----------------------------------------
# Network dataset
# ----------------------------------------
# Handle multiple GPUs
gpu_num = torch.cuda.device_count()
print("There are %d GPUs used" % gpu_num)
opt.train_batch_size *= gpu_num
#opt.val_batch_size *= gpu_num
opt.num_workers *= gpu_num
# Define the dataset
train_imglist = utils.get_jpgs(os.path.join(opt.in_path_train))
val_imglist = utils.get_jpgs(os.path.join(opt.in_path_val))
train_dataset = dataset.Qbayer2RGB_dataset(opt, 'train', train_imglist)
val_dataset = dataset.Qbayer2RGB_dataset(opt, 'val', val_imglist)
print('The overall number of training images:', len(train_imglist))
print('The overall number of validation images:', len(val_imglist))
# Define the dataloader
train_loader = torch.utils.data.DataLoader(train_dataset,
batch_size=opt.train_batch_size,
shuffle=True,
num_workers=opt.num_workers,
pin_memory=True)
val_loader = torch.utils.data.DataLoader(val_dataset,
batch_size=opt.val_batch_size,
shuffle=False,
num_workers=opt.num_workers,
pin_memory=True)
# ----------------------------------------
# Network training parameters
# ----------------------------------------
# Loss functions
criterion_L1 = torch.nn.L1Loss().cuda()
class ColorLoss(nn.Module):
def __init__(self):
super(ColorLoss, self).__init__()
self.L1loss = nn.L1Loss()
def RGB2YUV(self, RGB):
YUV = RGB.clone()
YUV[:,
0, :, :] = 0.299 * RGB[:,
0, :, :] + 0.587 * RGB[:,
1, :, :] + 0.114 * RGB[:,
2, :, :]
YUV[:,
1, :, :] = -0.14713 * RGB[:,
0, :, :] - 0.28886 * RGB[:,
1, :, :] + 0.436 * RGB[:,
2, :, :]
YUV[:,
2, :, :] = 0.615 * RGB[:,
0, :, :] - 0.51499 * RGB[:,
1, :, :] - 0.10001 * RGB[:,
2, :, :]
return YUV
def forward(self, x, y):
yuv_x = self.RGB2YUV(x)
yuv_y = self.RGB2YUV(y)
return self.L1loss(yuv_x, yuv_y)
yuv_loss = ColorLoss()
# Optimizers
optimizer_G = torch.optim.Adam(generator.parameters(),
lr=opt.lr_g,
betas=(opt.b1, opt.b2),
weight_decay=opt.weight_decay)
optimizer_D = torch.optim.Adam(generator.parameters(),
lr=opt.lr_d,
betas=(opt.b1, opt.b2),
weight_decay=opt.weight_decay)
# Learning rate decrease
def adjust_learning_rate(opt, epoch, iteration, optimizer, lr_gd):
# Set the learning rate to the initial LR decayed by "lr_decrease_factor" every "lr_decrease_epoch" epochs
if opt.lr_decrease_mode == 'epoch':
lr = lr_gd * (opt.lr_decrease_factor
**(epoch // opt.lr_decrease_epoch))
for param_group in optimizer.param_groups:
param_group['lr'] = lr
if opt.lr_decrease_mode == 'iter':
lr = lr_gd * (opt.lr_decrease_factor
**(iteration // opt.lr_decrease_iter))
for param_group in optimizer.param_groups:
param_group['lr'] = lr
# Save the model if pre_train == True
def save_model(opt, epoch, iteration, len_dataset, generator):
# Define the name of trained model
if opt.save_mode == 'epoch':
model_name = '%s_gan_noise%.3f_epoch%d_bs%d.pth' % (
opt.net_mode, opt.noise_level, epoch, opt.train_batch_size)
if opt.save_mode == 'iter':
model_name = '%s_gan_noise%.3f_iter%d_bs%d.pth' % (
opt.net_mode, opt.noise_level, iteration, opt.train_batch_size)
save_model_path = os.path.join(opt.save_path, opt.task_name,
model_name)
# Save model
if opt.multi_gpu == True:
if opt.save_mode == 'epoch':
if (epoch % opt.save_by_epoch
== 0) and (iteration % len_dataset == 0):
torch.save(generator.module.state_dict(), save_model_path)
print(
'The trained model is successfully saved at epoch %d' %
(epoch))
if opt.save_mode == 'iter':
if iteration % opt.save_by_iter == 0:
torch.save(generator.module.state_dict(), save_model_path)
print(
'The trained model is successfully saved at iteration %d'
% (iteration))
else:
if opt.save_mode == 'epoch':
if (epoch % opt.save_by_epoch
== 0) and (iteration % len_dataset == 0):
torch.save(generator.state_dict(), save_model_path)
print(
'The trained model is successfully saved at epoch %d' %
(epoch))
if opt.save_mode == 'iter':
if iteration % opt.save_by_iter == 0:
torch.save(generator.state_dict(), save_model_path)
print(
'The trained model is successfully saved at iteration %d'
% (iteration))
# ----------------------------------------
# Training
# ----------------------------------------
# Count start time
prev_time = time.time()
# Tensorboard
writer = SummaryWriter()
# For loop training
for epoch in range(opt.epochs):
# Record learning rate
for param_group in optimizer_G.param_groups:
writer.add_scalar('data/lr', param_group['lr'], epoch)
print('learning rate = ', param_group['lr'])
if epoch == 0:
iters_done = 0
### Training
for i, (in_img, RGBout_img) in enumerate(train_loader):
# To device
# A is for input image, B is for target image
in_img = in_img.cuda()
RGBout_img = RGBout_img.cuda()
## Train Discriminator
# Forward propagation
out = generator(in_img)
optimizer_D.zero_grad()
# Fake samples
fake_scalar_d = discriminator(in_img, out.detach())
true_scalar_d = discriminator(in_img, RGBout_img)
# Overall Loss and optimize
loss_D = -torch.mean(true_scalar_d) + torch.mean(fake_scalar_d)
loss_D.backward()
#torch.nn.utils.clip_grad_norm(discriminator.parameters(), opt.grad_clip_norm)
optimizer_D.step()
## Train Generator
# Forward propagation
out = generator(in_img)
# GAN loss
fake_scalar = discriminator(in_img, out)
L_gan = -torch.mean(fake_scalar) * opt.lambda_gan
# Perceptual loss features
fake_B_fea = perceptualnet(utils.normalize_ImageNet_stats(out))
true_B_fea = perceptualnet(
utils.normalize_ImageNet_stats(RGBout_img))
L_percep = opt.lambda_percep * criterion_L1(fake_B_fea, true_B_fea)
# Pixel loss
L_pixel = opt.lambda_pixel * criterion_L1(out, RGBout_img)
# Color loss
L_color = opt.lambda_color * yuv_loss(out, RGBout_img)
# Sum up to total loss
loss = L_pixel + L_percep + L_gan + L_color
# Record losses
writer.add_scalar('data/L_pixel', L_pixel.item(), iters_done)
writer.add_scalar('data/L_percep', L_percep.item(), iters_done)
writer.add_scalar('data/L_color', L_color.item(), iters_done)
writer.add_scalar('data/L_gan', L_gan.item(), iters_done)
writer.add_scalar('data/L_total', loss.item(), iters_done)
writer.add_scalar('data/loss_D', loss_D.item(), iters_done)
# Backpropagate gradients
optimizer_G.zero_grad()
loss.backward()
#torch.nn.utils.clip_grad_norm(generator.parameters(), opt.grad_clip_norm)
optimizer_G.step()
# Determine approximate time left
iters_done = epoch * len(train_loader) + i + 1
iters_left = opt.epochs * len(train_loader) - iters_done
time_left = datetime.timedelta(seconds=iters_left *
(time.time() - prev_time))
prev_time = time.time()
# Print log
print(
"\r[Epoch %d/%d] [Batch %d/%d] [Total Loss: %.4f] [L_pixel: %.4f]"
% ((epoch + 1), opt.epochs, i, len(train_loader), loss.item(),
L_pixel.item()))
print(
"\r[L_percep: %.4f] [L_color: %.4f] [L_gan: %.4f] [loss_D: %.4f] Time_left: %s"
% (L_percep.item(), L_color.item(), L_gan.item(),
loss_D.item(), time_left))
# Save model at certain epochs or iterations
save_model(opt, (epoch + 1), iters_done, len(train_loader),
generator)
# Learning rate decrease at certain epochs
adjust_learning_rate(opt, (epoch + 1), iters_done, optimizer_G,
opt.lr_g)
adjust_learning_rate(opt, (epoch + 1), iters_done, optimizer_D,
opt.lr_d)
### Sample data every epoch
if (epoch + 1) % 1 == 0:
img_list = [out, RGBout_img]
name_list = ['pred', 'gt']
utils.save_sample_png(sample_folder=sample_folder,
sample_name='train_epoch%d' % (epoch + 1),
img_list=img_list,
name_list=name_list,
pixel_max_cnt=255)
### Validation
val_PSNR = 0
num_of_val_image = 0
for j, (in_img, RGBout_img) in enumerate(val_loader):
# To device
# A is for input image, B is for target image
in_img = in_img.cuda()
RGBout_img = RGBout_img.cuda()
# Forward propagation
with torch.no_grad():
out = generator(in_img)
# Accumulate num of image and val_PSNR
num_of_val_image += in_img.shape[0]
val_PSNR += utils.psnr(out, RGBout_img, 1) * in_img.shape[0]
val_PSNR = val_PSNR / num_of_val_image
### Sample data every epoch
if (epoch + 1) % 1 == 0:
img_list = [out, RGBout_img]
name_list = ['pred', 'gt']
utils.save_sample_png(sample_folder=sample_folder,
sample_name='val_epoch%d' % (epoch + 1),
img_list=img_list,
name_list=name_list,
pixel_max_cnt=255)
# Record average PSNR
writer.add_scalar('data/val_PSNR', val_PSNR, epoch)
print('PSNR at epoch %d: %.4f' % ((epoch + 1), val_PSNR))
writer.close()
def Continue_train_WGAN(opt):
# ----------------------------------------
# Network training parameters
# ----------------------------------------
# cudnn benchmark
cudnn.benchmark = opt.cudnn_benchmark
# Loss functions
criterion_L1 = torch.nn.L1Loss().cuda()
# Initialize Generator
generator = utils.create_generator(opt)
discriminator = utils.create_discriminator(opt)
perceptualnet = utils.create_perceptualnet()
# To device
if opt.multi_gpu:
generator = nn.DataParallel(generator)
generator = generator.cuda()
discriminator = nn.DataParallel(discriminator)
discriminator = discriminator.cuda()
perceptualnet = nn.DataParallel(perceptualnet)
perceptualnet = perceptualnet.cuda()
else:
generator = generator.cuda()
discriminator = discriminator.cuda()
perceptualnet = perceptualnet.cuda()
# Optimizers
optimizer_G = torch.optim.Adam(generator.parameters(),
lr=opt.lr_g,
betas=(opt.b1, opt.b2),
weight_decay=opt.weight_decay)
optimizer_D = torch.optim.Adam(discriminator.parameters(),
lr=opt.lr_d,
betas=(opt.b1, opt.b2))
# Learning rate decrease
def adjust_learning_rate(opt, epoch, iteration, optimizer):
#Set the learning rate to the initial LR decayed by "lr_decrease_factor" every "lr_decrease_epoch" epochs
if opt.lr_decrease_mode == 'epoch':
lr = opt.lr_g * (opt.lr_decrease_factor
**(epoch // opt.lr_decrease_epoch))
for param_group in optimizer.param_groups:
param_group['lr'] = lr
if opt.lr_decrease_mode == 'iter':
lr = opt.lr_g * (opt.lr_decrease_factor
**(iteration // opt.lr_decrease_iter))
for param_group in optimizer.param_groups:
param_group['lr'] = lr
# Save the model if pre_train == True
def save_model(opt, epoch, iteration, len_dataset, generator):
"""Save the model at "checkpoint_interval" and its multiple"""
if opt.multi_gpu == True:
if opt.save_mode == 'epoch':
if (epoch % opt.save_by_epoch
== 0) and (iteration % len_dataset == 0):
if opt.save_name_mode:
torch.save(
generator.module, 'WGAN_%s_epoch%d_bs%d.pth' %
(opt.task, epoch, opt.batch_size))
print(
'The trained model is successfully saved at epoch %d'
% (epoch))
if opt.save_mode == 'iter':
if iteration % opt.save_by_iter == 0:
if opt.save_name_mode:
torch.save(
generator.module, 'WGAN_%s_iter%d_bs%d.pth' %
(opt.task, iteration, opt.batch_size))
print(
'The trained model is successfully saved at iteration %d'
% (iteration))
else:
if opt.save_mode == 'epoch':
if (epoch % opt.save_by_epoch
== 0) and (iteration % len_dataset == 0):
if opt.save_name_mode:
torch.save(
generator, 'WGAN_%s_epoch%d_bs%d.pth' %
(opt.task, epoch, opt.batch_size))
print(
'The trained model is successfully saved at epoch %d'
% (epoch))
if opt.save_mode == 'iter':
if iteration % opt.save_by_iter == 0:
if opt.save_name_mode:
torch.save(
generator, 'WGAN_%s_iter%d_bs%d.pth' %
(opt.task, iteration, opt.batch_size))
print(
'The trained model is successfully saved at iteration %d'
% (iteration))
# ----------------------------------------
# Network dataset
# ----------------------------------------
# Define the dataset
trainset = dataset.RAW2RGBDataset(opt)
print('The overall number of images:', len(trainset))
# Define the dataloader
dataloader = DataLoader(trainset,
batch_size=opt.batch_size,
shuffle=True,
num_workers=opt.num_workers,
pin_memory=True)
# ----------------------------------------
# Training
# ----------------------------------------
# Count start time
prev_time = time.time()
# For loop training
for epoch in range(opt.epochs):
for i, (true_input, true_target, true_sal) in enumerate(dataloader):
# To device
true_input = true_input.cuda()
true_target = true_target.cuda()
true_sal = true_sal.cuda()
true_sal3 = torch.cat((true_sal, true_sal, true_sal), 1).cuda()
# Train Discriminator
for j in range(opt.additional_training_d):
optimizer_D.zero_grad()
# Generator output
fake_target, fake_sal = generator(true_input)
# Fake samples
fake_block1, fake_block2, fake_block3, fake_scalar = discriminator(
true_input, fake_target.detach())
# True samples
true_block1, true_block2, true_block3, true_scalar = discriminator(
true_input, true_target)
'''
# Feature Matching Loss
FM_Loss = criterion_L1(fake_block1, true_block1) + criterion_L1(fake_block2, true_block2) + criterion_L1(fake_block3, true_block3)
'''
# Overall Loss and optimize
loss_D = -torch.mean(true_scalar) + torch.mean(fake_scalar)
loss_D.backward()
# Train Generator
optimizer_G.zero_grad()
fake_target, fake_sal = generator(true_input)
# L1 Loss
Pixellevel_L1_Loss = criterion_L1(fake_target, true_target)
# Attention Loss
true_Attn_target = true_target.mul(true_sal3)
fake_sal3 = torch.cat((fake_sal, fake_sal, fake_sal), 1)
fake_Attn_target = fake_target.mul(fake_sal3)
Attention_Loss = criterion_L1(fake_Attn_target, true_Attn_target)
# GAN Loss
fake_block1, fake_block2, fake_block3, fake_scalar = discriminator(
true_input, fake_target)
GAN_Loss = -torch.mean(fake_scalar)
# Perceptual Loss
fake_target = fake_target * 0.5 + 0.5
true_target = true_target * 0.5 + 0.5
fake_target = utils.normalize_ImageNet_stats(fake_target)
true_target = utils.normalize_ImageNet_stats(true_target)
fake_percep_feature = perceptualnet(fake_target)
true_percep_feature = perceptualnet(true_target)
Perceptual_Loss = criterion_L1(fake_percep_feature,
true_percep_feature)
# Overall Loss and optimize
loss = Pixellevel_L1_Loss + opt.lambda_attn * Attention_Loss + opt.lambda_gan * GAN_Loss + opt.lambda_percep * Perceptual_Loss
loss.backward()
optimizer_G.step()
# Determine approximate time left
iters_done = epoch * len(dataloader) + i
iters_left = opt.epochs * len(dataloader) - iters_done
time_left = datetime.timedelta(seconds=iters_left *
(time.time() - prev_time))
prev_time = time.time()
# Print log
print(
"\r[Epoch %d/%d] [Batch %d/%d] [Pixellevel L1 Loss: %.4f] [G Loss: %.4f] [D Loss: %.4f]"
% ((epoch + 1), opt.epochs, i, len(dataloader),
Pixellevel_L1_Loss.item(), GAN_Loss.item(), loss_D.item()))
print(
"\r[Attention Loss: %.4f] [Perceptual Loss: %.4f] Time_left: %s"
% (Attention_Loss.item(), Perceptual_Loss.item(), time_left))
# Save model at certain epochs or iterations
save_model(opt, (epoch + 1), (iters_done + 1), len(dataloader),
generator)
# Learning rate decrease at certain epochs
adjust_learning_rate(opt, (epoch + 1), (iters_done + 1),
optimizer_G)
adjust_learning_rate(opt, (epoch + 1), (iters_done + 1),
optimizer_D)
def WGAN_trainer(opt):
# ----------------------------------------
# Initialize training parameters
# ----------------------------------------
logger = Logger(opt)
checkpoint = restore(opt)
# cudnn benchmark accelerates the network
if opt.cudnn_benchmark == True:
cudnn.benchmark = True
else:
cudnn.benchmark = False
# --------------------------------------
# Initialize models
# --------------------------------------
generator, discriminator, perceptualnet = create_networks(opt, checkpoint)
# Loss functions
L1Loss = nn.L1Loss()
#FeatureMatchingLoss = FML1Loss(opt.fm_param)
# Optimizers
optimizer_g, optimizer_d = create_optimizers(generator, discriminator, opt,
checkpoint)
# Log metrics with wandb
wandb.watch(generator)
wandb.config.update(opt)
auto_sync_checkpoints_to_wandb()
# ----------------------------------------
# Initialize training dataset
# ----------------------------------------
# Define the dataset
trainset = dataset.InpaintDataset(opt)
print('The overall number of images equals to %d' % len(trainset))
# Define the dataloader
dataloader = DataLoader(trainset,
batch_size=opt.batch_size,
shuffle=True,
num_workers=opt.num_workers,
pin_memory=True)
# ----------------------------------------
# Training and Testing
# ----------------------------------------
# Initialize start time
prev_time = time.time()
initial_epoch = checkpoint['epoch'] if opt.restore else 0
n_iter = checkpoint['n_iter'] if opt.restore else 0
# training loop
for epoch in range(initial_epoch, opt.epochs):
for batch_idx, (img, mask) in enumerate(dataloader):
n_iter += 1
logger.begin(n_iter)
# Load mask (shape: [B, 1, H, W]), masked_img (shape: [B, 3, H, W]), img (shape: [B, 3, H, W]) and put it to cuda
img = img.cuda()
mask = mask.cuda()
### Train discriminator
optimizer_d.zero_grad()
# Generator output
first_out, second_out = generator(img, mask)
# forward propagation
first_out_wholeimg = img * (
1 - mask) + first_out * mask # in range [-1, 1]
second_out_wholeimg = img * (
1 - mask) + second_out * mask # in range [-1, 1]
if n_iter % opt.log_every == 1:
logger.add_image(img, 'image/training')
logger.add_image(mask, 'mask/training')
logger.add_image(first_out_wholeimg, 'image/first iteration')
logger.add_image(second_out_wholeimg, 'image/second iteration')
# Fake samples
fake_scalar = discriminator(second_out_wholeimg.detach(), mask)
# True samples
true_scalar = discriminator(img, mask)
# Overall Loss and optimize
loss_D = -torch.mean(true_scalar) + torch.mean(fake_scalar)
loss_D.backward()
optimizer_d.step()
### Train Generator
optimizer_g.zero_grad()
# Mask L1 Loss
first_MaskL1Loss = L1Loss(first_out_wholeimg, img)
second_MaskL1Loss = L1Loss(second_out_wholeimg, img)
# GAN Loss
fake_scalar = discriminator(second_out_wholeimg, mask)
GAN_Loss = -torch.mean(fake_scalar)
# Get the deep semantic feature maps, and compute Perceptual Loss
img = (img + 1) / 2 # in range [0, 1]
img = utils.normalize_ImageNet_stats(img) # in range of ImageNet
img_featuremaps = perceptualnet(img) # feature maps
second_out_wholeimg = (second_out_wholeimg +
1) / 2 # in range [0, 1]
second_out_wholeimg = utils.normalize_ImageNet_stats(
second_out_wholeimg)
second_out_wholeimg_featuremaps = perceptualnet(
second_out_wholeimg)
second_PerceptualLoss = L1Loss(second_out_wholeimg_featuremaps,
img_featuremaps)
# Compute losses
loss = first_MaskL1Loss + second_MaskL1Loss + opt.perceptual_param * second_PerceptualLoss + opt.gan_param * GAN_Loss
loss.backward()
optimizer_g.step()
# Determine approximate time left
batches_done = n_iter
batches_left = opt.epochs * len(dataloader) - batches_done
time_left = datetime.timedelta(seconds=batches_left *
(time.time() - prev_time))
prev_time = time.time()
logger.add_scalars({
'Epoch':
epoch + 1,
'Iteration':
n_iter,
'loss/first Mask L1 Loss':
first_MaskL1Loss.item(),
'loss/second Mask L1 Loss':
second_MaskL1Loss.item(),
'gan/D Loss':
loss_D.item(),
'gan/G Loss':
GAN_Loss.item(),
'Perceptual Loss':
second_PerceptualLoss.item()
})
# Print log
if n_iter % opt.log_every == 1:
print("\r[Epoch %d/%d] [Batch %d/%d] iteration %d" %
((epoch + 1), opt.epochs, batch_idx, len(dataloader),
n_iter))
print(
"\r[D Loss: %.5f] [G Loss: %.5f] [Perceptual Loss: %.5f] time_left: %s"
% (loss_D.item(), GAN_Loss.item(),
second_PerceptualLoss.item(), time_left))
if n_iter % opt.checkpoint_every == 1:
save_state(epoch=epoch,
batch=batch_idx,
n_iter=n_iter,
G=generator,
optimizer_g=optimizer_g,
D=discriminator,
optimizer_d=optimizer_d,
loss=loss,
opt=opt)
# Learning rate decrease
adjust_learning_rate(opt.lr_g, optimizer_g, (epoch + 1), opt)
adjust_learning_rate(opt.lr_d, optimizer_d, (epoch + 1), opt)
