""" # First update `train_config.py` to set paths to your dataset locations. # You may want to change `--num-workers` according to your machine's memory. # The default num-workers=8 may cause dataloader to exit unexpectedly when # machine is out of memory. # Stage 1 python train.py \ --model-variant mobilenetv3 \ --dataset videomatte \ --resolution-lr 512 \ --seq-length-lr 15 \ --learning-rate-backbone 0.0001 \ --learning-rate-aspp 0.0002 \ --learning-rate-decoder 0.0002 \ --learning-rate-refiner 0 \ --checkpoint-dir checkpoint/stage1 \ --log-dir log/stage1 \ --epoch-start 0 \ --epoch-end 20 # Stage 2 python train.py \ --model-variant mobilenetv3 \ --dataset videomatte \ --resolution-lr 512 \ --seq-length-lr 50 \ --learning-rate-backbone 0.00005 \ --learning-rate-aspp 0.0001 \ --learning-rate-decoder 0.0001 \ --learning-rate-refiner 0 \ --checkpoint checkpoint/stage1/epoch-19.pth \ --checkpoint-dir checkpoint/stage2 \ --log-dir log/stage2 \ --epoch-start 20 \ --epoch-end 22 # Stage 3 python train.py \ --model-variant mobilenetv3 \ --dataset videomatte \ --train-hr \ --resolution-lr 512 \ --resolution-hr 2048 \ --seq-length-lr 40 \ --seq-length-hr 6 \ --learning-rate-backbone 0.00001 \ --learning-rate-aspp 0.00001 \ --learning-rate-decoder 0.00001 \ --learning-rate-refiner 0.0002 \ --checkpoint checkpoint/stage2/epoch-21.pth \ --checkpoint-dir checkpoint/stage3 \ --log-dir log/stage3 \ --epoch-start 22 \ --epoch-end 23 # Stage 4 python train.py \ --model-variant mobilenetv3 \ --dataset imagematte \ --train-hr \ --resolution-lr 512 \ --resolution-hr 2048 \ --seq-length-lr 40 \ --seq-length-hr 6 \ --learning-rate-backbone 0.00001 \ --learning-rate-aspp 0.00001 \ --learning-rate-decoder 0.00005 \ --learning-rate-refiner 0.0002 \ --checkpoint checkpoint/stage3/epoch-22.pth \ --checkpoint-dir checkpoint/stage4 \ --log-dir log/stage4 \ --epoch-start 23 \ --epoch-end 28 """ import argparse import torch import random import os from torch import nn from torch import distributed as dist from torch import multiprocessing as mp from torch.nn import functional as F from torch.nn.parallel import DistributedDataParallel as DDP from torch.optim import Adam from torch.cuda.amp import autocast, GradScaler from torch.utils.data import DataLoader, ConcatDataset from torch.utils.data.distributed import DistributedSampler from torch.utils.tensorboard import SummaryWriter from torchvision.utils import make_grid from torchvision.transforms.functional import center_crop from tqdm import tqdm from dataset.videomatte import ( VideoMatteDataset, VideoMatteTrainAugmentation, VideoMatteValidAugmentation, ) from dataset.imagematte import ( ImageMatteDataset, ImageMatteAugmentation ) from dataset.coco import ( CocoPanopticDataset, CocoPanopticTrainAugmentation, ) from dataset.spd import ( SuperviselyPersonDataset ) from dataset.youtubevis import ( YouTubeVISDataset, YouTubeVISAugmentation ) from dataset.augmentation import ( TrainFrameSampler, ValidFrameSampler ) from model import MattingNetwork from train_config import DATA_PATHS from train_loss import matting_loss, segmentation_loss class Trainer: def __init__(self, rank, world_size): self.parse_args() self.init_distributed(rank, world_size) self.init_datasets() self.init_model() self.init_writer() self.train() self.cleanup() def parse_args(self): parser = argparse.ArgumentParser() # Model parser.add_argument('--model-variant', type=str, required=True, choices=['mobilenetv3', 'resnet50']) # Matting dataset parser.add_argument('--dataset', type=str, required=True, choices=['videomatte', 'imagematte']) # Learning rate parser.add_argument('--learning-rate-backbone', type=float, required=True) parser.add_argument('--learning-rate-aspp', type=float, required=True) parser.add_argument('--learning-rate-decoder', type=float, required=True) parser.add_argument('--learning-rate-refiner', type=float, required=True) # Training setting parser.add_argument('--train-hr', action='store_true') parser.add_argument('--resolution-lr', type=int, default=512) parser.add_argument('--resolution-hr', type=int, default=2048) parser.add_argument('--seq-length-lr', type=int, required=True) parser.add_argument('--seq-length-hr', type=int, default=6) parser.add_argument('--downsample-ratio', type=float, default=0.25) parser.add_argument('--batch-size-per-gpu', type=int, default=1) parser.add_argument('--num-workers', type=int, default=8) parser.add_argument('--epoch-start', type=int, default=0) parser.add_argument('--epoch-end', type=int, default=16) # Tensorboard logging parser.add_argument('--log-dir', type=str, required=True) parser.add_argument('--log-train-loss-interval', type=int, default=20) parser.add_argument('--log-train-images-interval', type=int, default=500) # Checkpoint loading and saving parser.add_argument('--checkpoint', type=str) parser.add_argument('--checkpoint-dir', type=str, required=True) parser.add_argument('--checkpoint-save-interval', type=int, default=500) # Distributed parser.add_argument('--distributed-addr', type=str, default='localhost') parser.add_argument('--distributed-port', type=str, default='12355') # Debugging parser.add_argument('--disable-progress-bar', action='store_true') parser.add_argument('--disable-validation', action='store_true') parser.add_argument('--disable-mixed-precision', action='store_true') self.args = parser.parse_args() def init_distributed(self, rank, world_size): self.rank = rank self.world_size = world_size self.log('Initializing distributed') os.environ['MASTER_ADDR'] = self.args.distributed_addr os.environ['MASTER_PORT'] = self.args.distributed_port dist.init_process_group("nccl", rank=rank, world_size=world_size) def init_datasets(self): self.log('Initializing matting datasets') size_hr = (self.args.resolution_hr, self.args.resolution_hr) size_lr = (self.args.resolution_lr, self.args.resolution_lr) # Matting datasets: if self.args.dataset == 'videomatte': self.dataset_lr_train = VideoMatteDataset( videomatte_dir=DATA_PATHS['videomatte']['train'], background_image_dir=DATA_PATHS['background_images']['train'], background_video_dir=DATA_PATHS['background_videos']['train'], size=self.args.resolution_lr, seq_length=self.args.seq_length_lr, seq_sampler=TrainFrameSampler(), transform=VideoMatteTrainAugmentation(size_lr)) if self.args.train_hr: self.dataset_hr_train = VideoMatteDataset( videomatte_dir=DATA_PATHS['videomatte']['train'], background_image_dir=DATA_PATHS['background_images']['train'], background_video_dir=DATA_PATHS['background_videos']['train'], size=self.args.resolution_hr, seq_length=self.args.seq_length_hr, seq_sampler=TrainFrameSampler(), transform=VideoMatteTrainAugmentation(size_hr)) self.dataset_valid = VideoMatteDataset( videomatte_dir=DATA_PATHS['videomatte']['valid'], background_image_dir=DATA_PATHS['background_images']['valid'], background_video_dir=DATA_PATHS['background_videos']['valid'], size=self.args.resolution_hr if self.args.train_hr else self.args.resolution_lr, seq_length=self.args.seq_length_hr if self.args.train_hr else self.args.seq_length_lr, seq_sampler=ValidFrameSampler(), transform=VideoMatteValidAugmentation(size_hr if self.args.train_hr else size_lr)) else: self.dataset_lr_train = ImageMatteDataset( imagematte_dir=DATA_PATHS['imagematte']['train'], background_image_dir=DATA_PATHS['background_images']['train'], background_video_dir=DATA_PATHS['background_videos']['train'], size=self.args.resolution_lr, seq_length=self.args.seq_length_lr, seq_sampler=TrainFrameSampler(), transform=ImageMatteAugmentation(size_lr)) if self.args.train_hr: self.dataset_hr_train = ImageMatteDataset( imagematte_dir=DATA_PATHS['imagematte']['train'], background_image_dir=DATA_PATHS['background_images']['train'], background_video_dir=DATA_PATHS['background_videos']['train'], size=self.args.resolution_hr, seq_length=self.args.seq_length_hr, seq_sampler=TrainFrameSampler(), transform=ImageMatteAugmentation(size_hr)) self.dataset_valid = ImageMatteDataset( imagematte_dir=DATA_PATHS['imagematte']['valid'], background_image_dir=DATA_PATHS['background_images']['valid'], background_video_dir=DATA_PATHS['background_videos']['valid'], size=self.args.resolution_hr if self.args.train_hr else self.args.resolution_lr, seq_length=self.args.seq_length_hr if self.args.train_hr else self.args.seq_length_lr, seq_sampler=ValidFrameSampler(), transform=ImageMatteAugmentation(size_hr if self.args.train_hr else size_lr)) # Matting dataloaders: self.datasampler_lr_train = DistributedSampler( dataset=self.dataset_lr_train, rank=self.rank, num_replicas=self.world_size, shuffle=True) self.dataloader_lr_train = DataLoader( dataset=self.dataset_lr_train, batch_size=self.args.batch_size_per_gpu, num_workers=self.args.num_workers, sampler=self.datasampler_lr_train, pin_memory=True) if self.args.train_hr: self.datasampler_hr_train = DistributedSampler( dataset=self.dataset_hr_train, rank=self.rank, num_replicas=self.world_size, shuffle=True) self.dataloader_hr_train = DataLoader( dataset=self.dataset_hr_train, batch_size=self.args.batch_size_per_gpu, num_workers=self.args.num_workers, sampler=self.datasampler_hr_train, pin_memory=True) self.dataloader_valid = DataLoader( dataset=self.dataset_valid, batch_size=self.args.batch_size_per_gpu, num_workers=self.args.num_workers, pin_memory=True) # Segementation datasets self.log('Initializing image segmentation datasets') self.dataset_seg_image = ConcatDataset([ CocoPanopticDataset( imgdir=DATA_PATHS['coco_panoptic']['imgdir'], anndir=DATA_PATHS['coco_panoptic']['anndir'], annfile=DATA_PATHS['coco_panoptic']['annfile'], transform=CocoPanopticTrainAugmentation(size_lr)), SuperviselyPersonDataset( imgdir=DATA_PATHS['spd']['imgdir'], segdir=DATA_PATHS['spd']['segdir'], transform=CocoPanopticTrainAugmentation(size_lr)) ]) self.datasampler_seg_image = DistributedSampler( dataset=self.dataset_seg_image, rank=self.rank, num_replicas=self.world_size, shuffle=True) self.dataloader_seg_image = DataLoader( dataset=self.dataset_seg_image, batch_size=self.args.batch_size_per_gpu * self.args.seq_length_lr, num_workers=self.args.num_workers, sampler=self.datasampler_seg_image, pin_memory=True) self.log('Initializing video segmentation datasets') self.dataset_seg_video = YouTubeVISDataset( videodir=DATA_PATHS['youtubevis']['videodir'], annfile=DATA_PATHS['youtubevis']['annfile'], size=self.args.resolution_lr, seq_length=self.args.seq_length_lr, seq_sampler=TrainFrameSampler(speed=[1]), transform=YouTubeVISAugmentation(size_lr)) self.datasampler_seg_video = DistributedSampler( dataset=self.dataset_seg_video, rank=self.rank, num_replicas=self.world_size, shuffle=True) self.dataloader_seg_video = DataLoader( dataset=self.dataset_seg_video, batch_size=self.args.batch_size_per_gpu, num_workers=self.args.num_workers, sampler=self.datasampler_seg_video, pin_memory=True) def init_model(self): self.log('Initializing model') self.model = MattingNetwork(self.args.model_variant, pretrained_backbone=True).to(self.rank) if self.args.checkpoint: self.log(f'Restoring from checkpoint: {self.args.checkpoint}') self.log(self.model.load_state_dict( torch.load(self.args.checkpoint, map_location=f'cuda:{self.rank}'))) self.model = nn.SyncBatchNorm.convert_sync_batchnorm(self.model) self.model_ddp = DDP(self.model, device_ids=[self.rank], broadcast_buffers=False, find_unused_parameters=True) self.optimizer = Adam([ {'params': self.model.backbone.parameters(), 'lr': self.args.learning_rate_backbone}, {'params': self.model.aspp.parameters(), 'lr': self.args.learning_rate_aspp}, {'params': self.model.decoder.parameters(), 'lr': self.args.learning_rate_decoder}, {'params': self.model.refiner.parameters(), 'lr': self.args.learning_rate_refiner}, ]) self.scaler = GradScaler() def init_writer(self): if self.rank == 0: self.log('Initializing writer') self.writer = SummaryWriter(self.args.log_dir) def train(self): for epoch in range(self.args.epoch_start, self.args.epoch_end): self.epoch = epoch self.step = epoch * len(self.dataloader_lr_train) if not self.args.disable_validation: self.validate() self.log(f'Training epoch: {epoch}') for true_fgr, true_pha, true_bgr in tqdm(self.dataloader_lr_train, disable=self.args.disable_progress_bar, dynamic_ncols=True): # Low resolution pass self.train_mat(true_fgr, true_pha, true_bgr, downsample_ratio=1, tag='lr') # High resolution pass if self.args.train_hr: true_fgr, true_pha, true_bgr = self.load_next_mat_hr_sample() self.train_mat(true_fgr, true_pha, true_bgr, downsample_ratio=self.args.downsample_ratio, tag='hr') # Segmentation pass if self.step % 2 == 0: true_img, true_seg = self.load_next_seg_video_sample() self.train_seg(true_img, true_seg, log_label='seg_video') else: true_img, true_seg = self.load_next_seg_image_sample() self.train_seg(true_img.unsqueeze(1), true_seg.unsqueeze(1), log_label='seg_image') if self.step % self.args.checkpoint_save_interval == 0: self.save() self.step += 1 def train_mat(self, true_fgr, true_pha, true_bgr, downsample_ratio, tag): true_fgr = true_fgr.to(self.rank, non_blocking=True) true_pha = true_pha.to(self.rank, non_blocking=True) true_bgr = true_bgr.to(self.rank, non_blocking=True) true_fgr, true_pha, true_bgr = self.random_crop(true_fgr, true_pha, true_bgr) true_src = true_fgr * true_pha + true_bgr * (1 - true_pha) with autocast(enabled=not self.args.disable_mixed_precision): pred_fgr, pred_pha = self.model_ddp(true_src, downsample_ratio=downsample_ratio)[:2] loss = matting_loss(pred_fgr, pred_pha, true_fgr, true_pha) self.scaler.scale(loss['total']).backward() self.scaler.step(self.optimizer) self.scaler.update() self.optimizer.zero_grad() if self.rank == 0 and self.step % self.args.log_train_loss_interval == 0: for loss_name, loss_value in loss.items(): self.writer.add_scalar(f'train_{tag}_{loss_name}', loss_value, self.step) if self.rank == 0 and self.step % self.args.log_train_images_interval == 0: self.writer.add_image(f'train_{tag}_pred_fgr', make_grid(pred_fgr.flatten(0, 1), nrow=pred_fgr.size(1)), self.step) self.writer.add_image(f'train_{tag}_pred_pha', make_grid(pred_pha.flatten(0, 1), nrow=pred_pha.size(1)), self.step) self.writer.add_image(f'train_{tag}_true_fgr', make_grid(true_fgr.flatten(0, 1), nrow=true_fgr.size(1)), self.step) self.writer.add_image(f'train_{tag}_true_pha', make_grid(true_pha.flatten(0, 1), nrow=true_pha.size(1)), self.step) self.writer.add_image(f'train_{tag}_true_src', make_grid(true_src.flatten(0, 1), nrow=true_src.size(1)), self.step) def train_seg(self, true_img, true_seg, log_label): true_img = true_img.to(self.rank, non_blocking=True) true_seg = true_seg.to(self.rank, non_blocking=True) true_img, true_seg = self.random_crop(true_img, true_seg) with autocast(enabled=not self.args.disable_mixed_precision): pred_seg = self.model_ddp(true_img, segmentation_pass=True)[0] loss = segmentation_loss(pred_seg, true_seg) self.scaler.scale(loss).backward() self.scaler.step(self.optimizer) self.scaler.update() self.optimizer.zero_grad() if self.rank == 0 and (self.step - self.step % 2) % self.args.log_train_loss_interval == 0: self.writer.add_scalar(f'{log_label}_loss', loss, self.step) if self.rank == 0 and (self.step - self.step % 2) % self.args.log_train_images_interval == 0: self.writer.add_image(f'{log_label}_pred_seg', make_grid(pred_seg.flatten(0, 1).float().sigmoid(), nrow=self.args.seq_length_lr), self.step) self.writer.add_image(f'{log_label}_true_seg', make_grid(true_seg.flatten(0, 1), nrow=self.args.seq_length_lr), self.step) self.writer.add_image(f'{log_label}_true_img', make_grid(true_img.flatten(0, 1), nrow=self.args.seq_length_lr), self.step) def load_next_mat_hr_sample(self): try: sample = next(self.dataiterator_mat_hr) except: self.datasampler_hr_train.set_epoch(self.datasampler_hr_train.epoch + 1) self.dataiterator_mat_hr = iter(self.dataloader_hr_train) sample = next(self.dataiterator_mat_hr) return sample def load_next_seg_video_sample(self): try: sample = next(self.dataiterator_seg_video) except: self.datasampler_seg_video.set_epoch(self.datasampler_seg_video.epoch + 1) self.dataiterator_seg_video = iter(self.dataloader_seg_video) sample = next(self.dataiterator_seg_video) return sample def load_next_seg_image_sample(self): try: sample = next(self.dataiterator_seg_image) except: self.datasampler_seg_image.set_epoch(self.datasampler_seg_image.epoch + 1) self.dataiterator_seg_image = iter(self.dataloader_seg_image) sample = next(self.dataiterator_seg_image) return sample def validate(self): if self.rank == 0: self.log(f'Validating at the start of epoch: {self.epoch}') self.model_ddp.eval() total_loss, total_count = 0, 0 with torch.no_grad(): with autocast(enabled=not self.args.disable_mixed_precision): for true_fgr, true_pha, true_bgr in tqdm(self.dataloader_valid, disable=self.args.disable_progress_bar, dynamic_ncols=True): true_fgr = true_fgr.to(self.rank, non_blocking=True) true_pha = true_pha.to(self.rank, non_blocking=True) true_bgr = true_bgr.to(self.rank, non_blocking=True) true_src = true_fgr * true_pha + true_bgr * (1 - true_pha) batch_size = true_src.size(0) pred_fgr, pred_pha = self.model(true_src)[:2] total_loss += matting_loss(pred_fgr, pred_pha, true_fgr, true_pha)['total'].item() * batch_size total_count += batch_size avg_loss = total_loss / total_count self.log(f'Validation set average loss: {avg_loss}') self.writer.add_scalar('valid_loss', avg_loss, self.step) self.model_ddp.train() dist.barrier() def random_crop(self, *imgs): h, w = imgs[0].shape[-2:] w = random.choice(range(w // 2, w)) h = random.choice(range(w // 2, h)) results = [] for img in imgs: B, T = img.shape[:2] img = img.flatten(0, 1) img = F.interpolate(img, (max(h, w), max(h, w)), mode='bilinear', align_corners=False) img = center_crop(img, (h, w)) img = img.reshape(B, T, *img.shape[1:]) results.append(img) return results def save(self): if self.rank == 0: os.makedirs(self.args.checkpoint_dir, exist_ok=True) torch.save(self.model.state_dict(), os.path.join(self.args.checkpoint_dir, f'epoch-{self.epoch}.pth')) self.log('Model saved') dist.barrier() def cleanup(self): dist.destroy_process_group() def log(self, msg): print(f'[GPU{self.rank}] {msg}') if __name__ == '__main__': world_size = torch.cuda.device_count() mp.spawn( Trainer, nprocs=world_size, args=(world_size,), join=True)