import os
import glob
import logging
import importlib
from tqdm import tqdm
import torch
import torch.nn as nn
import torch.nn.functional as F
from core.prefetch_dataloader import PrefetchDataLoader, CPUPrefetcher
from torch.utils.data.distributed import DistributedSampler
from torch.nn.parallel import DistributedDataParallel as DDP
import torchvision
from torch.utils.tensorboard import SummaryWriter
from core.lr_scheduler import MultiStepRestartLR, CosineAnnealingRestartLR
from core.loss import AdversarialLoss, PerceptualLoss, LPIPSLoss
from core.dataset import TrainDataset
from model.modules.flow_comp_raft import RAFT_bi, FlowLoss, EdgeLoss
from model.recurrent_flow_completion import RecurrentFlowCompleteNet
from RAFT.utils.flow_viz_pt import flow_to_image
class Trainer:
def __init__(self, config):
self.config = config
self.epoch = 0
self.iteration = 0
self.num_local_frames = config['train_data_loader']['num_local_frames']
self.num_ref_frames = config['train_data_loader']['num_ref_frames']
# setup data set and data loader
self.train_dataset = TrainDataset(config['train_data_loader'])
self.train_sampler = None
self.train_args = config['trainer']
if config['distributed']:
self.train_sampler = DistributedSampler(
self.train_dataset,
num_replicas=config['world_size'],
rank=config['global_rank'])
dataloader_args = dict(
dataset=self.train_dataset,
batch_size=self.train_args['batch_size'] // config['world_size'],
shuffle=(self.train_sampler is None),
num_workers=self.train_args['num_workers'],
sampler=self.train_sampler,
drop_last=True)
self.train_loader = PrefetchDataLoader(self.train_args['num_prefetch_queue'], **dataloader_args)
self.prefetcher = CPUPrefetcher(self.train_loader)
# set loss functions
self.adversarial_loss = AdversarialLoss(type=self.config['losses']['GAN_LOSS'])
self.adversarial_loss = self.adversarial_loss.to(self.config['device'])
self.l1_loss = nn.L1Loss()
# self.perc_loss = PerceptualLoss(
# layer_weights={'conv3_4': 0.25, 'conv4_4': 0.25, 'conv5_4': 0.5},
# use_input_norm=True,
# range_norm=True,
# criterion='l1'
# ).to(self.config['device'])
if self.config['losses']['perceptual_weight'] > 0:
self.perc_loss = LPIPSLoss(use_input_norm=True, range_norm=True).to(self.config['device'])
# self.flow_comp_loss = FlowCompletionLoss().to(self.config['device'])
# self.flow_comp_loss = FlowCompletionLoss(self.config['device'])
# set raft
self.fix_raft = RAFT_bi(device = self.config['device'])
self.fix_flow_complete = RecurrentFlowCompleteNet('/mnt/lustre/sczhou/VQGANs/CodeMOVI/experiments_model/recurrent_flow_completion_v5_train_flowcomp_v5/gen_760000.pth')
for p in self.fix_flow_complete.parameters():
p.requires_grad = False
self.fix_flow_complete.to(self.config['device'])
self.fix_flow_complete.eval()
# self.flow_loss = FlowLoss()
# setup models including generator and discriminator
net = importlib.import_module('model.' + config['model']['net'])
self.netG = net.InpaintGenerator()
# print(self.netG)
self.netG = self.netG.to(self.config['device'])
if not self.config['model'].get('no_dis', False):
if self.config['model'].get('dis_2d', False):
self.netD = net.Discriminator_2D(
in_channels=3,
use_sigmoid=config['losses']['GAN_LOSS'] != 'hinge')
else:
self.netD = net.Discriminator(
in_channels=3,
use_sigmoid=config['losses']['GAN_LOSS'] != 'hinge')
self.netD = self.netD.to(self.config['device'])
self.interp_mode = self.config['model']['interp_mode']
# setup optimizers and schedulers
self.setup_optimizers()
self.setup_schedulers()
self.load()
if config['distributed']:
self.netG = DDP(self.netG,
device_ids=[self.config['local_rank']],
output_device=self.config['local_rank'],
broadcast_buffers=True,
find_unused_parameters=True)
if not self.config['model']['no_dis']:
self.netD = DDP(self.netD,
device_ids=[self.config['local_rank']],
output_device=self.config['local_rank'],
broadcast_buffers=True,
find_unused_parameters=False)
# set summary writer
self.dis_writer = None
self.gen_writer = None
self.summary = {}
if self.config['global_rank'] == 0 or (not config['distributed']):
if not self.config['model']['no_dis']:
self.dis_writer = SummaryWriter(
os.path.join(config['save_dir'], 'dis'))
self.gen_writer = SummaryWriter(
os.path.join(config['save_dir'], 'gen'))
def setup_optimizers(self):
"""Set up optimizers."""
backbone_params = []
for name, param in self.netG.named_parameters():
if param.requires_grad:
backbone_params.append(param)
else:
print(f'Params {name} will not be optimized.')
optim_params = [
{
'params': backbone_params,
'lr': self.config['trainer']['lr']
},
]
self.optimG = torch.optim.Adam(optim_params,
betas=(self.config['trainer']['beta1'],
self.config['trainer']['beta2']))
if not self.config['model']['no_dis']:
self.optimD = torch.optim.Adam(
self.netD.parameters(),
lr=self.config['trainer']['lr'],
betas=(self.config['trainer']['beta1'],
self.config['trainer']['beta2']))
def setup_schedulers(self):
"""Set up schedulers."""
scheduler_opt = self.config['trainer']['scheduler']
scheduler_type = scheduler_opt.pop('type')
if scheduler_type in ['MultiStepLR', 'MultiStepRestartLR']:
self.scheG = MultiStepRestartLR(
self.optimG,
milestones=scheduler_opt['milestones'],
gamma=scheduler_opt['gamma'])
if not self.config['model']['no_dis']:
self.scheD = MultiStepRestartLR(
self.optimD,
milestones=scheduler_opt['milestones'],
gamma=scheduler_opt['gamma'])
elif scheduler_type == 'CosineAnnealingRestartLR':
self.scheG = CosineAnnealingRestartLR(
self.optimG,
periods=scheduler_opt['periods'],
restart_weights=scheduler_opt['restart_weights'],
eta_min=scheduler_opt['eta_min'])
if not self.config['model']['no_dis']:
self.scheD = CosineAnnealingRestartLR(
self.optimD,
periods=scheduler_opt['periods'],
restart_weights=scheduler_opt['restart_weights'],
eta_min=scheduler_opt['eta_min'])
else:
raise NotImplementedError(
f'Scheduler {scheduler_type} is not implemented yet.')
def update_learning_rate(self):
"""Update learning rate."""
self.scheG.step()
if not self.config['model']['no_dis']:
self.scheD.step()
def get_lr(self):
"""Get current learning rate."""
return self.optimG.param_groups[0]['lr']
def add_summary(self, writer, name, val):
"""Add tensorboard summary."""
if name not in self.summary:
self.summary[name] = 0
self.summary[name] += val
n = self.train_args['log_freq']
if writer is not None and self.iteration % n == 0:
writer.add_scalar(name, self.summary[name] / n, self.iteration)
self.summary[name] = 0
def load(self):
"""Load netG (and netD)."""
# get the latest checkpoint
model_path = self.config['save_dir']
# TODO: add resume name
if os.path.isfile(os.path.join(model_path, 'latest.ckpt')):
latest_epoch = open(os.path.join(model_path, 'latest.ckpt'),
'r').read().splitlines()[-1]
else:
ckpts = [
os.path.basename(i).split('.pth')[0]
for i in glob.glob(os.path.join(model_path, '*.pth'))
]
ckpts.sort()
latest_epoch = ckpts[-1][4:] if len(ckpts) > 0 else None
if latest_epoch is not None:
gen_path = os.path.join(model_path,
f'gen_{int(latest_epoch):06d}.pth')
dis_path = os.path.join(model_path,
f'dis_{int(latest_epoch):06d}.pth')
opt_path = os.path.join(model_path,
f'opt_{int(latest_epoch):06d}.pth')
if self.config['global_rank'] == 0:
print(f'Loading model from {gen_path}...')
dataG = torch.load(gen_path, map_location=self.config['device'])
self.netG.load_state_dict(dataG)
if not self.config['model']['no_dis'] and self.config['model']['load_d']:
dataD = torch.load(dis_path, map_location=self.config['device'])
self.netD.load_state_dict(dataD)
data_opt = torch.load(opt_path, map_location=self.config['device'])
self.optimG.load_state_dict(data_opt['optimG'])
# self.scheG.load_state_dict(data_opt['scheG'])
if not self.config['model']['no_dis'] and self.config['model']['load_d']:
self.optimD.load_state_dict(data_opt['optimD'])
# self.scheD.load_state_dict(data_opt['scheD'])
self.epoch = data_opt['epoch']
self.iteration = data_opt['iteration']
else:
gen_path = self.config['trainer'].get('gen_path', None)
dis_path = self.config['trainer'].get('dis_path', None)
opt_path = self.config['trainer'].get('opt_path', None)
if gen_path is not None:
if self.config['global_rank'] == 0:
print(f'Loading Gen-Net from {gen_path}...')
dataG = torch.load(gen_path, map_location=self.config['device'])
self.netG.load_state_dict(dataG)
if dis_path is not None and not self.config['model']['no_dis'] and self.config['model']['load_d']:
if self.config['global_rank'] == 0:
print(f'Loading Dis-Net from {dis_path}...')
dataD = torch.load(dis_path, map_location=self.config['device'])
self.netD.load_state_dict(dataD)
if opt_path is not None:
data_opt = torch.load(opt_path, map_location=self.config['device'])
self.optimG.load_state_dict(data_opt['optimG'])
self.scheG.load_state_dict(data_opt['scheG'])
if not self.config['model']['no_dis'] and self.config['model']['load_d']:
self.optimD.load_state_dict(data_opt['optimD'])
self.scheD.load_state_dict(data_opt['scheD'])
else:
if self.config['global_rank'] == 0:
print('Warnning: There is no trained model found.'
'An initialized model will be used.')
def save(self, it):
"""Save parameters every eval_epoch"""
if self.config['global_rank'] == 0:
# configure path
gen_path = os.path.join(self.config['save_dir'],
f'gen_{it:06d}.pth')
dis_path = os.path.join(self.config['save_dir'],
f'dis_{it:06d}.pth')
opt_path = os.path.join(self.config['save_dir'],
f'opt_{it:06d}.pth')
print(f'\nsaving model to {gen_path} ...')
# remove .module for saving
if isinstance(self.netG, torch.nn.DataParallel) or isinstance(self.netG, DDP):
netG = self.netG.module
if not self.config['model']['no_dis']:
netD = self.netD.module
else:
netG = self.netG
if not self.config['model']['no_dis']:
netD = self.netD
# save checkpoints
torch.save(netG.state_dict(), gen_path)
if not self.config['model']['no_dis']:
torch.save(netD.state_dict(), dis_path)
torch.save(
{
'epoch': self.epoch,
'iteration': self.iteration,
'optimG': self.optimG.state_dict(),
'optimD': self.optimD.state_dict(),
'scheG': self.scheG.state_dict(),
'scheD': self.scheD.state_dict()
}, opt_path)
else:
torch.save(
{
'epoch': self.epoch,
'iteration': self.iteration,
'optimG': self.optimG.state_dict(),
'scheG': self.scheG.state_dict()
}, opt_path)
latest_path = os.path.join(self.config['save_dir'], 'latest.ckpt')
os.system(f"echo {it:06d} > {latest_path}")
def train(self):
"""training entry"""
pbar = range(int(self.train_args['iterations']))
if self.config['global_rank'] == 0:
pbar = tqdm(pbar,
initial=self.iteration,
dynamic_ncols=True,
smoothing=0.01)
os.makedirs('logs', exist_ok=True)
logging.basicConfig(
level=logging.INFO,
format="%(asctime)s %(filename)s[line:%(lineno)d]"
"%(levelname)s %(message)s",
datefmt="%a, %d %b %Y %H:%M:%S",
filename=f"logs/{self.config['save_dir'].split('/')[-1]}.log",
filemode='w')
while True:
self.epoch += 1
self.prefetcher.reset()
if self.config['distributed']:
self.train_sampler.set_epoch(self.epoch)
self._train_epoch(pbar)
if self.iteration > self.train_args['iterations']:
break
print('\nEnd training....')
def _train_epoch(self, pbar):
"""Process input and calculate loss every training epoch"""
device = self.config['device']
train_data = self.prefetcher.next()
while train_data is not None:
self.iteration += 1
frames, masks, flows_f, flows_b, _ = train_data
frames, masks = frames.to(device), masks.to(device).float()
l_t = self.num_local_frames
b, t, c, h, w = frames.size()
gt_local_frames = frames[:, :l_t, ...]
local_masks = masks[:, :l_t, ...].contiguous()
masked_frames = frames * (1 - masks)
masked_local_frames = masked_frames[:, :l_t, ...]
# get gt optical flow
if flows_f[0] == 'None' or flows_b[0] == 'None':
gt_flows_bi = self.fix_raft(gt_local_frames)
else:
gt_flows_bi = (flows_f.to(device), flows_b.to(device))
# ---- complete flow ----
pred_flows_bi, _ = self.fix_flow_complete.forward_bidirect_flow(gt_flows_bi, local_masks)
pred_flows_bi = self.fix_flow_complete.combine_flow(gt_flows_bi, pred_flows_bi, local_masks)
# pred_flows_bi = gt_flows_bi
# ---- image propagation ----
prop_imgs, updated_local_masks = self.netG.module.img_propagation(masked_local_frames, pred_flows_bi, local_masks, interpolation=self.interp_mode)
updated_masks = masks.clone()
updated_masks[:, :l_t, ...] = updated_local_masks.view(b, l_t, 1, h, w)
updated_frames = masked_frames.clone()
prop_local_frames = gt_local_frames * (1-local_masks) + prop_imgs.view(b, l_t, 3, h, w) * local_masks # merge
updated_frames[:, :l_t, ...] = prop_local_frames
# ---- feature propagation + Transformer ----
pred_imgs = self.netG(updated_frames, pred_flows_bi, masks, updated_masks, l_t)
pred_imgs = pred_imgs.view(b, -1, c, h, w)
# get the local frames
pred_local_frames = pred_imgs[:, :l_t, ...]
comp_local_frames = gt_local_frames * (1. - local_masks) + pred_local_frames * local_masks
comp_imgs = frames * (1. - masks) + pred_imgs * masks
gen_loss = 0
dis_loss = 0
# optimize net_g
if not self.config['model']['no_dis']:
for p in self.netD.parameters():
p.requires_grad = False
self.optimG.zero_grad()
# generator l1 loss
hole_loss = self.l1_loss(pred_imgs * masks, frames * masks)
hole_loss = hole_loss / torch.mean(masks) * self.config['losses']['hole_weight']
gen_loss += hole_loss
self.add_summary(self.gen_writer, 'loss/hole_loss', hole_loss.item())
valid_loss = self.l1_loss(pred_imgs * (1 - masks), frames * (1 - masks))
valid_loss = valid_loss / torch.mean(1-masks) * self.config['losses']['valid_weight']
gen_loss += valid_loss
self.add_summary(self.gen_writer, 'loss/valid_loss', valid_loss.item())
# perceptual loss
if self.config['losses']['perceptual_weight'] > 0:
perc_loss = self.perc_loss(pred_imgs.view(-1,3,h,w), frames.view(-1,3,h,w))[0] * self.config['losses']['perceptual_weight']
gen_loss += perc_loss
self.add_summary(self.gen_writer, 'loss/perc_loss', perc_loss.item())
# gan loss
if not self.config['model']['no_dis']:
# generator adversarial loss
gen_clip = self.netD(comp_imgs)
gan_loss = self.adversarial_loss(gen_clip, True, False)
gan_loss = gan_loss * self.config['losses']['adversarial_weight']
gen_loss += gan_loss
self.add_summary(self.gen_writer, 'loss/gan_loss', gan_loss.item())
gen_loss.backward()
self.optimG.step()
if not self.config['model']['no_dis']:
# optimize net_d
for p in self.netD.parameters():
p.requires_grad = True
self.optimD.zero_grad()
# discriminator adversarial loss
real_clip = self.netD(frames)
fake_clip = self.netD(comp_imgs.detach())
dis_real_loss = self.adversarial_loss(real_clip, True, True)
dis_fake_loss = self.adversarial_loss(fake_clip, False, True)
dis_loss += (dis_real_loss + dis_fake_loss) / 2
self.add_summary(self.dis_writer, 'loss/dis_vid_real', dis_real_loss.item())
self.add_summary(self.dis_writer, 'loss/dis_vid_fake', dis_fake_loss.item())
dis_loss.backward()
self.optimD.step()
self.update_learning_rate()
# write image to tensorboard
if self.iteration % 200 == 0:
# img to cpu
t = 0
gt_local_frames_cpu = ((gt_local_frames.view(b,-1,3,h,w) + 1)/2.0).cpu()
masked_local_frames = ((masked_local_frames.view(b,-1,3,h,w) + 1)/2.0).cpu()
prop_local_frames_cpu = ((prop_local_frames.view(b,-1,3,h,w) + 1)/2.0).cpu()
pred_local_frames_cpu = ((pred_local_frames.view(b,-1,3,h,w) + 1)/2.0).cpu()
img_results = torch.cat([masked_local_frames[0][t], gt_local_frames_cpu[0][t],
prop_local_frames_cpu[0][t], pred_local_frames_cpu[0][t]], 1)
img_results = torchvision.utils.make_grid(img_results, nrow=1, normalize=True)
if self.gen_writer is not None:
self.gen_writer.add_image(f'img/img:inp-gt-res-{t}', img_results, self.iteration)
t = 5
if masked_local_frames.shape[1] > 5:
img_results = torch.cat([masked_local_frames[0][t], gt_local_frames_cpu[0][t],
prop_local_frames_cpu[0][t], pred_local_frames_cpu[0][t]], 1)
img_results = torchvision.utils.make_grid(img_results, nrow=1, normalize=True)
if self.gen_writer is not None:
self.gen_writer.add_image(f'img/img:inp-gt-res-{t}', img_results, self.iteration)
# flow to cpu
gt_flows_forward_cpu = flow_to_image(gt_flows_bi[0][0]).cpu()
masked_flows_forward_cpu = (gt_flows_forward_cpu[0] * (1-local_masks[0][0].cpu())).to(gt_flows_forward_cpu)
pred_flows_forward_cpu = flow_to_image(pred_flows_bi[0][0]).cpu()
flow_results = torch.cat([gt_flows_forward_cpu[0], masked_flows_forward_cpu, pred_flows_forward_cpu[0]], 1)
if self.gen_writer is not None:
self.gen_writer.add_image('img/flow:gt-pred', flow_results, self.iteration)
# console logs
if self.config['global_rank'] == 0:
pbar.update(1)
if not self.config['model']['no_dis']:
pbar.set_description((f"d: {dis_loss.item():.3f}; "
f"hole: {hole_loss.item():.3f}; "
f"valid: {valid_loss.item():.3f}"))
else:
pbar.set_description((f"hole: {hole_loss.item():.3f}; "
f"valid: {valid_loss.item():.3f}"))
if self.iteration % self.train_args['log_freq'] == 0:
if not self.config['model']['no_dis']:
logging.info(f"[Iter {self.iteration}] "
f"d: {dis_loss.item():.4f}; "
f"hole: {hole_loss.item():.4f}; "
f"valid: {valid_loss.item():.4f}")
else:
logging.info(f"[Iter {self.iteration}] "
f"hole: {hole_loss.item():.4f}; "
f"valid: {valid_loss.item():.4f}")
# saving models
if self.iteration % self.train_args['save_freq'] == 0:
self.save(int(self.iteration))
if self.iteration > self.train_args['iterations']:
break
train_data = self.prefetcher.next()