from __future__ import print_function import argparse import math import os import queue import time import cv2 import torch import torch.backends.cudnn as cudnn import torch.nn as nn import torch.optim as optim import torchvision.transforms as transform_lib from torch.utils.data import ConcatDataset, DataLoader, WeightedRandomSampler from torchvision.transforms import CenterCrop import lib.TrainTransforms as transforms from lib.videoloader import VideosDataset from lib.videoloader_imagenet import VideosDataset_ImageNet from models.ColorVidNet import ColorVidNet from models.ContextualLoss import ContextualLoss, ContextualLoss_forward from models.FrameColor import frame_colorization from models.GAN_models import Discriminator_x64 from models.NonlocalNet import (NonlocalWeightedAverage, VGG19_pytorch, WarpNet, WeightedAverage, WeightedAverage_color) from tensorboardX import SummaryWriter from utils.util import (batch_lab2rgb_transpose_mc, feature_normalize, l1_loss, mkdir_if_not, mse_loss, parse, tensor_lab2rgb, uncenter_l, weighted_l1_loss, weighted_mse_loss) from utils.util_distortion import (CenterPad_threshold, Normalize, RGB2Lab, ToTensor) from utils.util_tensorboard import TBImageRecorder, value_logger from utils.warping import WarpingLayer cv2.setNumThreads(0) parser = argparse.ArgumentParser() parser.add_argument("--data_root", default="/home/zhanbo/remote/video/video_pair3/", type=str) parser.add_argument("--data_root_imagenet", default="/home/zhanbo/remote/imagenet_ref_pair/", type=str) parser.add_argument("--gpu_ids", type=str, default="0,1,2,3", help="separate by comma") parser.add_argument("--workers", type=int, default=8) parser.add_argument("--batch_size", type=int, default=16) parser.add_argument("--image_size", type=int, default=[216, 384]) parser.add_argument("--ic", type=int, default=7) parser.add_argument("--epoch", type=int, default=40) parser.add_argument("--resume_epoch", type=int, default=0) parser.add_argument("--resume", type=bool, default=False) parser.add_argument("--load_pretrained_model", type=bool, default=True) parser.add_argument("--lr", type=float, default=0.00001) parser.add_argument("--beta1", type=float, default=0.5) parser.add_argument("--lr_step", type=int, default=100) parser.add_argument("--lr_gamma", type=float, default=0.1) parser.add_argument("--checkpoint_dir", type=str, default="checkpoints/default") parser.add_argument("--tb_log_step", type=int, default=50) parser.add_argument("--print_step", type=int, default=2) parser.add_argument("--real_reference_probability", type=float, default=0.7) parser.add_argument("--nonzero_placeholder_probability", type=float, default=0.0) parser.add_argument("--with_bad", type=bool, default=True) parser.add_argument("--with_mid", type=bool, default=True) parser.add_argument("--domain_invariant", type=bool, default=False) parser.add_argument("--weigth_l1", type=float, default=2.0) parser.add_argument("--weight_contextual", type=float, default="0.2") parser.add_argument("--weight_perceptual", type=float, default="0.001") parser.add_argument("--weight_smoothness", type=float, default="5.0") parser.add_argument("--weight_gan", type=float, default="0.2") parser.add_argument("--weight_nonlocal_smoothness", type=float, default="0.0") parser.add_argument("--weight_nonlocal_consistent", type=float, default="0.0") parser.add_argument("--weight_consistent", type=float, default="0.02") parser.add_argument("--luminance_noise", type=float, default="2.0") parser.add_argument("--permute_data", type=bool, default=True) parser.add_argument("--contextual_loss_direction", type=str, default="forward", help="forward or backward matching") def image_logger_fn( I_last_lab, I_current_lab, I_reference_lab, I_last_lab_predict, I_current_lab_predict, I_last_nonlocal_lab, I_current_nonlocal_lab, ): I_last_image = batch_lab2rgb_transpose_mc(I_last_lab[0:32, 0:1, :, :], I_last_lab[0:32, 1:3, :, :]) I_current_image = batch_lab2rgb_transpose_mc(I_current_lab[0:32, 0:1, :, :], I_current_lab[0:32, 1:3, :, :]) I_reference_image = batch_lab2rgb_transpose_mc(I_reference_lab[0:32, 0:1, :, :], I_reference_lab[0:32, 1:3, :, :]) I_last_image_predict = batch_lab2rgb_transpose_mc( I_last_lab_predict[0:32, 0:1, :, :], I_last_lab_predict[0:32, 1:3, :, :] ) I_current_image_predict = batch_lab2rgb_transpose_mc( I_current_lab_predict[0:32, 0:1, :, :], I_current_lab_predict[0:32, 1:3, :, :] ) I_last_nonlocal_image = batch_lab2rgb_transpose_mc( I_last_nonlocal_lab[0:32, 0:1, :, :], I_last_nonlocal_lab[0:32, 1:3, :, :] ) I_current_nonlocal_image = batch_lab2rgb_transpose_mc( I_current_nonlocal_lab[0:32, 0:1, :, :], I_current_nonlocal_lab[0:32, 1:3, :, :] ) img_info = {} img_info["1_I_last"] = I_last_image img_info["2_I_current"] = I_current_image img_info["3_I_reference"] = I_reference_image img_info["4_I_last_predict"] = I_last_image_predict img_info["5_I_curren_predict"] = I_current_image_predict img_info["6_I_last_nonlocal"] = I_last_nonlocal_image img_info["7_I_current_nonlocal"] = I_current_nonlocal_image return img_info def training_logger(): try: if total_iter % opt.print_step == 0: print("processing time:", elapsed) print( "Epoch %d, Step[%d/%d], lr: %f, total_loss: %.2f" % ( epoch, ((iter + 1) % iter_num_per_epoch), iter_num_per_epoch, step_optim_scheduler_g.get_last_lr()[0], total_loss.item(), ) ) value_logger( tb_writer, total_iter, loss_info={ "l1_loss": l1_loss.item(), "feat_loss": feat_loss.item(), "contextual_loss_total": contextual_loss_total.item(), "smoothness_loss": smoothness_loss.item(), "nonlocal_smoothness_loss": nonlocal_smoothness_loss.item(), "nonlocal_consistent_loss": nonlocal_consistent_loss.item(), "consistent_loss": consistent_loss.item(), "generator_loss": generator_loss.item(), "discriminator_loss": discriminator_loss.item(), "total_loss": total_loss.item(), }, ) if total_iter % opt.tb_log_step == 0: I_last_nonlocal_lab = torch.cat( (I_last_lab[:, 0:1, :, :], I_last_nonlocal_lab_predict[:, 1:3, :, :]), dim=1 ) I_current_nonlocal_lab = torch.cat( (I_current_lab[:, 0:1, :, :], I_current_nonlocal_lab_predict[:, 1:3, :, :]), dim=1 ) I_last_lab_predict = torch.cat((I_last_l, I_last_ab_predict), dim=1) data_queue.put( ( ( I_last_lab.cpu(), I_current_lab.cpu(), I_reference_lab.cpu(), I_last_lab_predict.cpu(), I_current_lab_predict.cpu(), I_last_nonlocal_lab.cpu(), I_current_nonlocal_lab, ), total_iter, ) ) if total_iter % 2000 == 0: if len(opt.gpu_ids) > 1: torch.save( nonlocal_net.module.state_dict(), os.path.join(opt.checkpoint_dir, "nonlocal_net_iter_%d.pth") % total_iter, ) torch.save( colornet.module.state_dict(), os.path.join(opt.checkpoint_dir, "colornet_iter_%d.pth") % total_iter, ) torch.save( discriminator.module.state_dict(), os.path.join(opt.checkpoint_dir, "discriminator_iter_%d.pth") % total_iter, ) else: torch.save( nonlocal_net.state_dict(), os.path.join(opt.checkpoint_dir, "nonlocal_net_iter_%d.pth") % total_iter, ) torch.save(colornet.state_dict(), os.path.join(opt.checkpoint_dir, "colornet_iter_%d.pth") % total_iter) torch.save( discriminator.state_dict(), os.path.join(opt.checkpoint_dir, "discriminator_iter_%d.pth") % total_iter, ) # save the state for resume if total_iter % 2000 == 0: print("saving the checkpoint") if len(opt.gpu_ids) > 1: state = { "total_iter": total_iter + 1, "epoch": epoch, "colornet_state": colornet.module.state_dict(), "nonlocal_net_state": nonlocal_net.module.state_dict(), "discriminator_state": discriminator.module.state_dict(), "optimizer_g": optimizer_g.state_dict(), "optimizer_d": optimizer_d.state_dict(), "optimizer_schedule_g": step_optim_scheduler_g.state_dict(), "optimizer_schedule_d": step_optim_scheduler_g.state_dict(), } else: state = { "total_iter": total_iter + 1, "epoch": epoch, "colornet_state": colornet.state_dict(), "nonlocal_net_state": nonlocal_net.state_dict(), "discriminator_state": discriminator.state_dict(), "optimizer_g": optimizer_g.state_dict(), "optimizer_d": optimizer_d.state_dict(), "optimizer_schedule_g": step_optim_scheduler_g.state_dict(), "optimizer_schedule_d": step_optim_scheduler_d.state_dict(), } torch.save(state, os.path.join(opt.checkpoint_dir, "learning_checkpoint.pth")) except Exception as e: print("Exception during output") print(e) def gpu_setup(): os.environ["CUDA_DEVICE_ORDER"] = "PCI_BUS_ID" cudnn.benchmark = True opt.gpu_ids = [int(x) for x in opt.gpu_ids.split(",")] torch.cuda.set_device(opt.gpu_ids[0]) device = torch.device("cuda") print("running on GPU", opt.gpu_ids) return device def load_data(): print("initializing dataloader") transforms_video = [ CenterCrop(opt.image_size), RGB2Lab(), ToTensor(), Normalize(), ] transforms_imagenet = [CenterPad_threshold(opt.image_size), RGB2Lab(), ToTensor(), Normalize()] extra_reference_transform = [ transform_lib.RandomHorizontalFlip(0.5), transform_lib.RandomResizedCrop(480, (0.98, 1.0), ratio=(0.8, 1.2)), ] train_dataset_video = VideosDataset( data_root=opt.data_root, epoch=opt.epoch, image_size=opt.image_size, image_transform=transforms.Compose(transforms_video), real_reference_probability=opt.real_reference_probability, nonzero_placeholder_probability=opt.nonzero_placeholder_probability, ) train_dataset_imagenet = VideosDataset_ImageNet( data_root=opt.data_root_imagenet, image_size=opt.image_size, epoch=opt.epoch, with_bad=opt.with_bad, with_mid=opt.with_mid, transforms_imagenet=transforms_imagenet, distortion_level=4, brightnessjitter=5, nonzero_placeholder_probability=opt.nonzero_placeholder_probability, extra_reference_transform=extra_reference_transform, real_reference_probability=opt.real_reference_probability, ) video_training_length = len(train_dataset_video) imagenet_training_length = len(train_dataset_imagenet) dataset_training_length = train_dataset_video.real_len + train_dataset_imagenet.real_len dataset_combined = ConcatDataset([train_dataset_video, train_dataset_imagenet]) sampler = WeightedRandomSampler( [1] * video_training_length + [1] * imagenet_training_length, dataset_training_length * opt.epoch ) data_loader = DataLoader( dataset_combined, batch_size=opt.batch_size, shuffle=False, num_workers=opt.workers, pin_memory=True, drop_last=True, sampler=sampler, ) return dataset_training_length, train_dataset_video, train_dataset_imagenet, data_loader def define_loss(): print("defining loss") ab_criterion = nn.SmoothL1Loss().to(device) nonlocal_criterion = nn.SmoothL1Loss().to(device) feat_l2_criterion = nn.MSELoss().to(device) feat_l1_criterion = nn.SmoothL1Loss().to(device) contextual_loss = ContextualLoss().to(device) contextual_forward_loss = ContextualLoss_forward().to(device) BCE_stable = nn.BCEWithLogitsLoss().to(device) return contextual_loss, contextual_forward_loss def define_optimizer(): print("defining optimizer") optimizer_g = optim.Adam( [{"params": nonlocal_net.parameters(), "lr": 1e-5}, {"params": colornet.parameters(), "lr": 2e-4}], betas=(0.5, 0.999), eps=1e-5, amsgrad=True, ) optimizer_d = optim.Adam( filter(lambda p: p.requires_grad, discriminator.parameters()), lr=2 * 1e-4, betas=(0.5, 0.999) ) return optimizer_g, optimizer_d def resume_model(): print("resuming the learning") checkpoint = torch.load(os.path.join(opt.checkpoint_dir, "learning_checkpoint.pth")) total_iter = checkpoint["total_iter"] epoch = checkpoint["epoch"] colornet.load_state_dict(checkpoint["colornet_state"]) nonlocal_net.load_state_dict(checkpoint["nonlocal_net_state"]) discriminator.load_state_dict(checkpoint["discriminator_state"]) optimizer_g.load_state_dict(checkpoint["optimizer_g"]) optimizer_d.load_state_dict(checkpoint["optimizer_d"]) step_optim_scheduler_g.load_state_dict(checkpoint["optimizer_schedule_g"]) step_optim_scheduler_d.load_state_dict(checkpoint["optimizer_schedule_d"]) def to_device( colornet, nonlocal_net, discriminator, vggnet, contextual_loss, contextual_forward_loss, weighted_layer_color, nonlocal_weighted_layer, warping_layer, instancenorm, ): print("moving models to device") colornet = torch.nn.DataParallel(colornet.to(device), device_ids=opt.gpu_ids) nonlocal_net = torch.nn.DataParallel(nonlocal_net.to(device), device_ids=opt.gpu_ids) discriminator = torch.nn.DataParallel(discriminator.to(device), device_ids=opt.gpu_ids) vggnet = torch.nn.DataParallel(vggnet.to(device), device_ids=opt.gpu_ids) contextual_loss = torch.nn.DataParallel(contextual_loss.to(device), device_ids=opt.gpu_ids) contextual_forward_loss = torch.nn.DataParallel(contextual_forward_loss.to(device), device_ids=opt.gpu_ids) weighted_layer_color = torch.nn.DataParallel(weighted_layer_color.to(device), device_ids=opt.gpu_ids) nonlocal_weighted_layer = torch.nn.DataParallel(nonlocal_weighted_layer.to(device), device_ids=opt.gpu_ids) warping_layer = torch.nn.DataParallel(warping_layer.to(device), device_ids=opt.gpu_ids) instancenorm = torch.nn.DataParallel(instancenorm.to(device), device_ids=opt.gpu_ids) return ( vggnet, nonlocal_net, colornet, discriminator, instancenorm, contextual_loss, contextual_forward_loss, weighted_layer_color, nonlocal_weighted_layer, warping_layer, ) def loss_init(): print("initializing losses") zero_loss = torch.Tensor([0]).to(device) ( feat_loss, contextual_loss_total, smoothness_loss, nonlocal_smoothness_loss, consistent_loss, nonlocal_consistent_loss, generator_loss, discriminator_loss, ) = (zero_loss, zero_loss, zero_loss, zero_loss, zero_loss, zero_loss, zero_loss, zero_loss) return ( feat_loss, contextual_loss_total, smoothness_loss, nonlocal_smoothness_loss, consistent_loss, nonlocal_consistent_loss, generator_loss, discriminator_loss, ) def video_colorization(): # colorization for the last frame I_last_ab_predict, I_last_nonlocal_lab_predict, features_last_gray = frame_colorization( I_last_lab, I_reference_lab, placeholder_lab, features_B, vggnet, nonlocal_net, colornet, feature_noise=0, luminance_noise=opt.luminance_noise, ) I_last_lab_predict = torch.cat((I_last_l, I_last_ab_predict), dim=1) # colorization for the current frame I_current_ab_predict, I_current_nonlocal_lab_predict, features_current_gray = frame_colorization( I_current_lab, I_reference_lab, I_last_lab_predict, features_B, vggnet, nonlocal_net, colornet, feature_noise=0, luminance_noise=opt.luminance_noise, ) I_current_lab_predict = torch.cat((I_last_l, I_current_ab_predict), dim=1) return I_current_ab_predict, I_last_ab_predict, I_current_nonlocal_lab_predict, I_last_nonlocal_lab_predict if __name__ == "__main__": torch.multiprocessing.set_start_method("spawn", force=True) opt = parse(parser) opt.data_root = opt.data_root.split(",")[0] opt.data_root_imagenet = opt.data_root_imagenet.split(",")[0] mkdir_if_not(opt.checkpoint_dir) mkdir_if_not("./runs/") device = gpu_setup() dataset_training_length, train_dataset_video, train_dataset_imagenet, data_loader = load_data() tb_writer = SummaryWriter() data_queue = queue.Queue() tb_image_reorder = TBImageRecorder(tb_writer, image_logger_fn, data_queue) tb_image_reorder.start() # define network nonlocal_net = WarpNet(opt.batch_size) colornet = ColorVidNet(opt.ic) discriminator = Discriminator_x64(ndf=64) colornet.to(device) nonlocal_net.to(device) discriminator.to(device) weighted_layer = WeightedAverage() weighted_layer_color = WeightedAverage_color() nonlocal_weighted_layer = NonlocalWeightedAverage() instancenorm = nn.InstanceNorm2d(512, affine=False) warping_layer = WarpingLayer("gpu") vggnet = VGG19_pytorch() vggnet.load_state_dict(torch.load("data/vgg19_conv.pth")) vggnet.eval() for param in vggnet.parameters(): param.requires_grad = False # load pre-trained model if opt.load_pretrained_model: nonlocal_pretain_path = os.path.join("checkpoints/video_moredata_l1/", "nonlocal_net_iter_76000.pth") nonlocal_net.load_state_dict(torch.load(nonlocal_pretain_path)) color_test_path = "checkpoints/video_moredata_l1/" + "colornet_iter_76000.pth" colornet.load_state_dict(torch.load(color_test_path)) # define loss function contextual_loss, contextual_forward_loss = define_loss() # define optimizer optimizer_g, optimizer_d = define_optimizer() step_optim_scheduler_g = optim.lr_scheduler.StepLR(optimizer_g, step_size=opt.lr_step, gamma=opt.lr_gamma) step_optim_scheduler_d = optim.lr_scheduler.StepLR(optimizer_d, step_size=opt.lr_step, gamma=opt.lr_gamma) # define others downsampling_by2 = nn.AvgPool2d(kernel_size=2).to(device) downsampling_by4 = nn.AvgPool2d(kernel_size=4).to(device) # dataset info iter_num_per_epoch = dataset_training_length // opt.batch_size total_iter = opt.resume_epoch * iter_num_per_epoch print( "train_dataset info, real_len: %d, epoch_len: %d, iter_num_per_epoch: %d" % (dataset_training_length, len(train_dataset_video) + len(train_dataset_imagenet), iter_num_per_epoch) ) if opt.resume: resume_model() # move to GPU processing device = opt.gpu_ids[0] ( vggnet, nonlocal_net, colornet, discriminator, instancenorm, contextual_loss, contextual_forward_loss, weighted_layer_color, nonlocal_weighted_layer, warping_layer, ) = to_device( colornet, nonlocal_net, discriminator, vggnet, contextual_loss, contextual_forward_loss, weighted_layer_color, nonlocal_weighted_layer, warping_layer, instancenorm, ) ( feat_loss, contextual_loss_total, smoothness_loss, nonlocal_smoothness_loss, consistent_loss, nonlocal_consistent_loss, generator_loss, discriminator_loss, ) = loss_init() # %% Training print("start training") for iter, data in enumerate(data_loader): start_time = time.time() total_iter += 1 epoch = math.ceil(total_iter / iter_num_per_epoch) ###### LOADING DATA SAMPLE ###### ( I_last_lab, I_current_lab, I_reference_lab, flow_forward, flow_backward, mask, placeholder_lab, self_ref_flag, ) = data I_last_lab = I_last_lab.to(device) I_current_lab = I_current_lab.to(device) I_reference_lab = I_reference_lab.to(device) flow_forward = flow_forward.to(device) flow_backward = flow_backward.to(device) mask = mask.to(device) placeholder_lab = placeholder_lab.to(device) self_ref_flag = self_ref_flag.to(device) I_last_l = I_last_lab[:, 0:1, :, :] I_last_ab = I_last_lab[:, 1:3, :, :] I_current_l = I_current_lab[:, 0:1, :, :] I_current_ab = I_current_lab[:, 1:3, :, :] I_reference_l = I_reference_lab[:, 0:1, :, :] I_reference_ab = I_reference_lab[:, 1:3, :, :] I_reference_rgb = tensor_lab2rgb(torch.cat((uncenter_l(I_reference_l), I_reference_ab), dim=1)) features_B = vggnet(I_reference_rgb, ["r12", "r22", "r32", "r42", "r52"], preprocess=True) ###### COLORIZATION ###### ( I_current_ab_predict, I_last_ab_predict, I_current_nonlocal_lab_predict, I_last_nonlocal_lab_predict, ) = video_colorization() ###### UPDATE DISCRIMINATOR ###### if opt.weight_gan > 0: optimizer_g.zero_grad() optimizer_d.zero_grad() fake_data_lab = torch.cat( (uncenter_l(I_current_l), I_current_ab_predict, uncenter_l(I_last_l), I_last_ab_predict), dim=1 ) real_data_lab = torch.cat((uncenter_l(I_current_l), I_current_ab, uncenter_l(I_last_l), I_last_ab), dim=1) if opt.permute_data: batch_index = torch.arange(-1, opt.batch_size - 1, dtype=torch.long) real_data_lab = real_data_lab[batch_index, ...] y_pred_fake, feature_pred_fake = discriminator(fake_data_lab.detach()) y_pred_real, feature_pred_real = discriminator(real_data_lab.detach()) y = torch.ones_like(y_pred_real) y2 = torch.zeros_like(y_pred_real) discriminator_loss = ( torch.mean((y_pred_real - torch.mean(y_pred_fake) - y) ** 2) + torch.mean((y_pred_fake - torch.mean(y_pred_real) + y) ** 2) ) / 2 discriminator_loss.backward() optimizer_d.step() ###### UPDATE GENERATOR ###### optimizer_g.zero_grad() optimizer_d.zero_grad() # extract vgg features for both output and original image I_predict_rgb = tensor_lab2rgb(torch.cat((uncenter_l(I_current_l), I_current_ab_predict), dim=1)) predict_relu1_1, predict_relu2_1, predict_relu3_1, predict_relu4_1, predict_relu5_1 = vggnet( I_predict_rgb, ["r12", "r22", "r32", "r42", "r52"], preprocess=True ) I_current_rgb = tensor_lab2rgb(torch.cat((uncenter_l(I_current_l), I_current_ab), dim=1)) A_relu1_1, A_relu2_1, A_relu3_1, A_relu4_1, A_relu5_1 = vggnet( I_current_rgb, ["r12", "r22", "r32", "r42", "r52"], preprocess=True ) B_relu1_1, B_relu2_1, B_relu3_1, B_relu4_1, B_relu5_1 = features_B ###### LOSS COMPUTE ###### # l1 loss if opt.weigth_l1 > 0: sample_weights = (self_ref_flag[:, 1:3, :, :]) / (sum(self_ref_flag[:, 0, 0, 0]) + 1e-5) l1_loss = weighted_l1_loss(I_current_ab_predict, I_current_ab, sample_weights) * opt.weigth_l1 # generator loss if opt.weight_gan > 0: y_pred_fake, feature_pred_fake = discriminator(fake_data_lab) y_pred_real, feature_pred_real = discriminator(real_data_lab) generator_loss = ( ( torch.mean((y_pred_real - torch.mean(y_pred_fake) + y) ** 2) + torch.mean((y_pred_fake - torch.mean(y_pred_real) - y) ** 2) ) / 2 * opt.weight_gan ) # feature loss if opt.domain_invariant: feat_loss = ( mse_loss(instancenorm(predict_relu5_1), instancenorm(A_relu5_1.detach())) * opt.weight_perceptual * 1e5 * 0.2 ) else: feat_loss = mse_loss(predict_relu5_1, A_relu5_1.detach()) * opt.weight_perceptual # contextual loss if opt.contextual_loss_direction == "backward": contextual_style5_1 = torch.mean(contextual_loss(predict_relu5_1, B_relu5_1.detach())) * 8 contextual_style4_1 = torch.mean(contextual_loss(predict_relu4_1, B_relu4_1.detach())) * 4 contextual_style3_1 = ( torch.mean(contextual_loss(downsampling_by2(predict_relu3_1), downsampling_by2(B_relu3_1.detach()))) * 2 ) else: contextual_style5_1 = torch.mean(contextual_forward_loss(predict_relu5_1, B_relu5_1.detach())) * 8 contextual_style4_1 = torch.mean(contextual_forward_loss(predict_relu4_1, B_relu4_1.detach())) * 4 contextual_style3_1 = ( torch.mean( contextual_forward_loss(downsampling_by2(predict_relu3_1), downsampling_by2(B_relu3_1.detach())) ) * 2 ) if opt.weight_contextual > 0: contextual_loss_total = ( contextual_style5_1 + contextual_style4_1 + contextual_style3_1 ) * opt.weight_contextual # smoothness loss if opt.weight_smoothness > 0: scale_factor = 1 I_current_lab_predict = torch.cat((I_current_l, I_current_ab_predict), dim=1) IA_ab_weighed = weighted_layer_color( I_current_lab, I_current_lab_predict, patch_size=3, alpha=10, scale_factor=scale_factor ) smoothness_loss = ( mse_loss(nn.functional.interpolate(I_current_ab_predict, scale_factor=scale_factor), IA_ab_weighed) * opt.weight_smoothness ) if opt.weight_nonlocal_smoothness > 0: scale_factor = 0.25 alpha_nonlocal_smoothness = 0.5 nonlocal_smooth_feature = feature_normalize(A_relu2_1) I_current_lab_predict = torch.cat((I_current_l, I_current_ab_predict), dim=1) I_current_ab_weighted_nonlocal = nonlocal_weighted_layer( I_current_lab_predict, nonlocal_smooth_feature.detach(), patch_size=3, alpha=alpha_nonlocal_smoothness, scale_factor=scale_factor, ) nonlocal_smoothness_loss = ( mse_loss( nn.functional.interpolate(I_current_ab_predict, scale_factor=scale_factor), I_current_ab_weighted_nonlocal, ) * opt.weight_nonlocal_smoothness ) if opt.weight_consistent: I_current_lab_predict_warp = warping_layer(I_current_lab_predict, flow_forward) I_current_ab_predict_warp = I_current_lab_predict_warp[:, 1:3, :, :] consistent_loss = ( weighted_mse_loss(I_current_ab_predict_warp, I_last_ab_predict, mask) * opt.weight_consistent ) if opt.weight_nonlocal_consistent: I_current_nonlocal_lab_predict_warp = warping_layer(I_current_nonlocal_lab_predict, flow_forward) nonlocal_consistent_loss = ( weighted_mse_loss( I_current_nonlocal_lab_predict_warp[:, 1:3, :, :], I_last_nonlocal_lab_predict[:, 1:3, :, :], mask ) * opt.weight_nonlocal_consistent ) # total loss total_loss = ( l1_loss + feat_loss + contextual_loss_total + smoothness_loss + nonlocal_smoothness_loss + consistent_loss + nonlocal_consistent_loss + generator_loss ) total_loss.backward() optimizer_g.step() end_time = time.time() elapsed = end_time - start_time training_logger() step_optim_scheduler_g.step() step_optim_scheduler_d.step() data_queue.put((None, None))