739 lines
29 KiB
Python
739 lines
29 KiB
Python
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))
|