Spaces:
Running
on
A10G
Running
on
A10G
| import torch | |
| from torch.utils.data import DataLoader | |
| from torch.utils.tensorboard import SummaryWriter | |
| import argparse | |
| import numpy as np | |
| import os | |
| from data import build_train_dataset | |
| from gmflow.gmflow import GMFlow | |
| from loss import flow_loss_func | |
| from evaluate import (validate_chairs, validate_things, validate_sintel, validate_kitti, | |
| create_sintel_submission, create_kitti_submission, inference_on_dir) | |
| from utils.logger import Logger | |
| from utils import misc | |
| from utils.dist_utils import get_dist_info, init_dist, setup_for_distributed | |
| def get_args_parser(): | |
| parser = argparse.ArgumentParser() | |
| # dataset | |
| parser.add_argument('--checkpoint_dir', default='tmp', type=str, | |
| help='where to save the training log and models') | |
| parser.add_argument('--stage', default='chairs', type=str, | |
| help='training stage') | |
| parser.add_argument('--image_size', default=[384, 512], type=int, nargs='+', | |
| help='image size for training') | |
| parser.add_argument('--padding_factor', default=16, type=int, | |
| help='the input should be divisible by padding_factor, otherwise do padding') | |
| parser.add_argument('--max_flow', default=400, type=int, | |
| help='exclude very large motions during training') | |
| parser.add_argument('--val_dataset', default=['chairs'], type=str, nargs='+', | |
| help='validation dataset') | |
| parser.add_argument('--with_speed_metric', action='store_true', | |
| help='with speed metric when evaluation') | |
| # training | |
| parser.add_argument('--lr', default=4e-4, type=float) | |
| parser.add_argument('--batch_size', default=12, type=int) | |
| parser.add_argument('--num_workers', default=4, type=int) | |
| parser.add_argument('--weight_decay', default=1e-4, type=float) | |
| parser.add_argument('--grad_clip', default=1.0, type=float) | |
| parser.add_argument('--num_steps', default=100000, type=int) | |
| parser.add_argument('--seed', default=326, type=int) | |
| parser.add_argument('--summary_freq', default=100, type=int) | |
| parser.add_argument('--val_freq', default=10000, type=int) | |
| parser.add_argument('--save_ckpt_freq', default=10000, type=int) | |
| parser.add_argument('--save_latest_ckpt_freq', default=1000, type=int) | |
| # resume pretrained model or resume training | |
| parser.add_argument('--resume', default=None, type=str, | |
| help='resume from pretrain model for finetuing or resume from terminated training') | |
| parser.add_argument('--strict_resume', action='store_true') | |
| parser.add_argument('--no_resume_optimizer', action='store_true') | |
| # GMFlow model | |
| parser.add_argument('--num_scales', default=1, type=int, | |
| help='basic gmflow model uses a single 1/8 feature, the refinement uses 1/4 feature') | |
| parser.add_argument('--feature_channels', default=128, type=int) | |
| parser.add_argument('--upsample_factor', default=8, type=int) | |
| parser.add_argument('--num_transformer_layers', default=6, type=int) | |
| parser.add_argument('--num_head', default=1, type=int) | |
| parser.add_argument('--attention_type', default='swin', type=str) | |
| parser.add_argument('--ffn_dim_expansion', default=4, type=int) | |
| parser.add_argument('--attn_splits_list', default=[2], type=int, nargs='+', | |
| help='number of splits in attention') | |
| parser.add_argument('--corr_radius_list', default=[-1], type=int, nargs='+', | |
| help='correlation radius for matching, -1 indicates global matching') | |
| parser.add_argument('--prop_radius_list', default=[-1], type=int, nargs='+', | |
| help='self-attention radius for flow propagation, -1 indicates global attention') | |
| # loss | |
| parser.add_argument('--gamma', default=0.9, type=float, | |
| help='loss weight') | |
| # evaluation | |
| parser.add_argument('--eval', action='store_true') | |
| parser.add_argument('--save_eval_to_file', action='store_true') | |
| parser.add_argument('--evaluate_matched_unmatched', action='store_true') | |
| # inference on a directory | |
| parser.add_argument('--inference_dir', default=None, type=str) | |
| parser.add_argument('--inference_size', default=None, type=int, nargs='+', | |
| help='can specify the inference size') | |
| parser.add_argument('--dir_paired_data', action='store_true', | |
| help='Paired data in a dir instead of a sequence') | |
| parser.add_argument('--save_flo_flow', action='store_true') | |
| parser.add_argument('--pred_bidir_flow', action='store_true', | |
| help='predict bidirectional flow') | |
| parser.add_argument('--fwd_bwd_consistency_check', action='store_true', | |
| help='forward backward consistency check with bidirection flow') | |
| # predict on sintel and kitti test set for submission | |
| parser.add_argument('--submission', action='store_true', | |
| help='submission to sintel or kitti test sets') | |
| parser.add_argument('--output_path', default='output', type=str, | |
| help='where to save the prediction results') | |
| parser.add_argument('--save_vis_flow', action='store_true', | |
| help='visualize flow prediction as .png image') | |
| parser.add_argument('--no_save_flo', action='store_true', | |
| help='not save flow as .flo') | |
| # distributed training | |
| parser.add_argument('--local_rank', default=0, type=int) | |
| parser.add_argument('--distributed', action='store_true') | |
| parser.add_argument('--launcher', default='none', type=str, choices=['none', 'pytorch']) | |
| parser.add_argument('--gpu_ids', default=0, type=int, nargs='+') | |
| parser.add_argument('--count_time', action='store_true', | |
| help='measure the inference time on sintel') | |
| return parser | |
| def main(args): | |
| if not args.eval and not args.submission and args.inference_dir is None: | |
| if args.local_rank == 0: | |
| print('pytorch version:', torch.__version__) | |
| print(args) | |
| misc.save_args(args) | |
| misc.check_path(args.checkpoint_dir) | |
| misc.save_command(args.checkpoint_dir) | |
| seed = args.seed | |
| torch.manual_seed(seed) | |
| np.random.seed(seed) | |
| torch.backends.cudnn.benchmark = True | |
| if args.launcher == 'none': | |
| args.distributed = False | |
| device = torch.device('cuda' if torch.cuda.is_available() else 'cpu') | |
| else: | |
| args.distributed = True | |
| # adjust batch size for each gpu | |
| assert args.batch_size % torch.cuda.device_count() == 0 | |
| args.batch_size = args.batch_size // torch.cuda.device_count() | |
| dist_params = dict(backend='nccl') | |
| init_dist(args.launcher, **dist_params) | |
| # re-set gpu_ids with distributed training mode | |
| _, world_size = get_dist_info() | |
| args.gpu_ids = range(world_size) | |
| device = torch.device('cuda:{}'.format(args.local_rank)) | |
| setup_for_distributed(args.local_rank == 0) | |
| # model | |
| model = GMFlow(feature_channels=args.feature_channels, | |
| num_scales=args.num_scales, | |
| upsample_factor=args.upsample_factor, | |
| num_head=args.num_head, | |
| attention_type=args.attention_type, | |
| ffn_dim_expansion=args.ffn_dim_expansion, | |
| num_transformer_layers=args.num_transformer_layers, | |
| ).to(device) | |
| if not args.eval and not args.submission and not args.inference_dir: | |
| print('Model definition:') | |
| print(model) | |
| if args.distributed: | |
| model = torch.nn.parallel.DistributedDataParallel( | |
| model.to(device), | |
| device_ids=[args.local_rank], | |
| output_device=args.local_rank) | |
| model_without_ddp = model.module | |
| else: | |
| if torch.cuda.device_count() > 1: | |
| print('Use %d GPUs' % torch.cuda.device_count()) | |
| model = torch.nn.DataParallel(model) | |
| model_without_ddp = model.module | |
| else: | |
| model_without_ddp = model | |
| num_params = sum(p.numel() for p in model.parameters()) | |
| print('Number of params:', num_params) | |
| if not args.eval and not args.submission and args.inference_dir is None: | |
| save_name = '%d_parameters' % num_params | |
| open(os.path.join(args.checkpoint_dir, save_name), 'a').close() | |
| optimizer = torch.optim.AdamW(model_without_ddp.parameters(), lr=args.lr, | |
| weight_decay=args.weight_decay) | |
| start_epoch = 0 | |
| start_step = 0 | |
| # resume checkpoints | |
| if args.resume: | |
| print('Load checkpoint: %s' % args.resume) | |
| loc = 'cuda:{}'.format(args.local_rank) | |
| checkpoint = torch.load(args.resume, map_location=loc) | |
| weights = checkpoint['model'] if 'model' in checkpoint else checkpoint | |
| model_without_ddp.load_state_dict(weights, strict=args.strict_resume) | |
| if 'optimizer' in checkpoint and 'step' in checkpoint and 'epoch' in checkpoint and not \ | |
| args.no_resume_optimizer: | |
| print('Load optimizer') | |
| optimizer.load_state_dict(checkpoint['optimizer']) | |
| start_epoch = checkpoint['epoch'] | |
| start_step = checkpoint['step'] | |
| print('start_epoch: %d, start_step: %d' % (start_epoch, start_step)) | |
| # evaluate | |
| if args.eval: | |
| val_results = {} | |
| if 'chairs' in args.val_dataset: | |
| results_dict = validate_chairs(model_without_ddp, | |
| with_speed_metric=args.with_speed_metric, | |
| attn_splits_list=args.attn_splits_list, | |
| corr_radius_list=args.corr_radius_list, | |
| prop_radius_list=args.prop_radius_list, | |
| ) | |
| val_results.update(results_dict) | |
| if 'things' in args.val_dataset: | |
| results_dict = validate_things(model_without_ddp, | |
| padding_factor=args.padding_factor, | |
| with_speed_metric=args.with_speed_metric, | |
| attn_splits_list=args.attn_splits_list, | |
| corr_radius_list=args.corr_radius_list, | |
| prop_radius_list=args.prop_radius_list, | |
| ) | |
| val_results.update(results_dict) | |
| if 'sintel' in args.val_dataset: | |
| results_dict = validate_sintel(model_without_ddp, | |
| count_time=args.count_time, | |
| padding_factor=args.padding_factor, | |
| with_speed_metric=args.with_speed_metric, | |
| evaluate_matched_unmatched=args.evaluate_matched_unmatched, | |
| attn_splits_list=args.attn_splits_list, | |
| corr_radius_list=args.corr_radius_list, | |
| prop_radius_list=args.prop_radius_list, | |
| ) | |
| val_results.update(results_dict) | |
| if 'kitti' in args.val_dataset: | |
| results_dict = validate_kitti(model_without_ddp, | |
| padding_factor=args.padding_factor, | |
| with_speed_metric=args.with_speed_metric, | |
| attn_splits_list=args.attn_splits_list, | |
| corr_radius_list=args.corr_radius_list, | |
| prop_radius_list=args.prop_radius_list, | |
| ) | |
| val_results.update(results_dict) | |
| if args.save_eval_to_file: | |
| misc.check_path(args.checkpoint_dir) | |
| val_file = os.path.join(args.checkpoint_dir, 'val_results.txt') | |
| with open(val_file, 'a') as f: | |
| f.write('\neval results after training done\n\n') | |
| metrics = ['chairs_epe', 'chairs_s0_10', 'chairs_s10_40', 'chairs_s40+', | |
| 'things_clean_epe', 'things_clean_s0_10', 'things_clean_s10_40', 'things_clean_s40+', | |
| 'things_final_epe', 'things_final_s0_10', 'things_final_s10_40', 'things_final_s40+', | |
| 'sintel_clean_epe', 'sintel_clean_s0_10', 'sintel_clean_s10_40', 'sintel_clean_s40+', | |
| 'sintel_final_epe', 'sintel_final_s0_10', 'sintel_final_s10_40', 'sintel_final_s40+', | |
| 'kitti_epe', 'kitti_f1', 'kitti_s0_10', 'kitti_s10_40', 'kitti_s40+', | |
| ] | |
| eval_metrics = [] | |
| for metric in metrics: | |
| if metric in val_results.keys(): | |
| eval_metrics.append(metric) | |
| metrics_values = [val_results[metric] for metric in eval_metrics] | |
| num_metrics = len(eval_metrics) | |
| # save as markdown format | |
| f.write(("| {:>20} " * num_metrics + '\n').format(*eval_metrics)) | |
| f.write(("| {:20.3f} " * num_metrics).format(*metrics_values)) | |
| f.write('\n\n') | |
| return | |
| # Sintel and KITTI submission | |
| if args.submission: | |
| # NOTE: args.val_dataset is a list | |
| if args.val_dataset[0] == 'sintel': | |
| create_sintel_submission(model_without_ddp, | |
| output_path=args.output_path, | |
| padding_factor=args.padding_factor, | |
| save_vis_flow=args.save_vis_flow, | |
| no_save_flo=args.no_save_flo, | |
| attn_splits_list=args.attn_splits_list, | |
| corr_radius_list=args.corr_radius_list, | |
| prop_radius_list=args.prop_radius_list, | |
| ) | |
| elif args.val_dataset[0] == 'kitti': | |
| create_kitti_submission(model_without_ddp, | |
| output_path=args.output_path, | |
| padding_factor=args.padding_factor, | |
| save_vis_flow=args.save_vis_flow, | |
| attn_splits_list=args.attn_splits_list, | |
| corr_radius_list=args.corr_radius_list, | |
| prop_radius_list=args.prop_radius_list, | |
| ) | |
| else: | |
| raise ValueError(f'Not supported dataset for submission') | |
| return | |
| # inferece on a dir | |
| if args.inference_dir is not None: | |
| inference_on_dir(model_without_ddp, | |
| inference_dir=args.inference_dir, | |
| output_path=args.output_path, | |
| padding_factor=args.padding_factor, | |
| inference_size=args.inference_size, | |
| paired_data=args.dir_paired_data, | |
| save_flo_flow=args.save_flo_flow, | |
| attn_splits_list=args.attn_splits_list, | |
| corr_radius_list=args.corr_radius_list, | |
| prop_radius_list=args.prop_radius_list, | |
| pred_bidir_flow=args.pred_bidir_flow, | |
| fwd_bwd_consistency_check=args.fwd_bwd_consistency_check, | |
| ) | |
| return | |
| # training datset | |
| train_dataset = build_train_dataset(args) | |
| print('Number of training images:', len(train_dataset)) | |
| # Multi-processing | |
| if args.distributed: | |
| train_sampler = torch.utils.data.distributed.DistributedSampler( | |
| train_dataset, | |
| num_replicas=torch.cuda.device_count(), | |
| rank=args.local_rank) | |
| else: | |
| train_sampler = None | |
| shuffle = False if args.distributed else True | |
| train_loader = torch.utils.data.DataLoader(train_dataset, batch_size=args.batch_size, | |
| shuffle=shuffle, num_workers=args.num_workers, | |
| pin_memory=True, drop_last=True, | |
| sampler=train_sampler) | |
| last_epoch = start_step if args.resume and start_step > 0 else -1 | |
| lr_scheduler = torch.optim.lr_scheduler.OneCycleLR( | |
| optimizer, args.lr, | |
| args.num_steps + 10, | |
| pct_start=0.05, | |
| cycle_momentum=False, | |
| anneal_strategy='cos', | |
| last_epoch=last_epoch, | |
| ) | |
| if args.local_rank == 0: | |
| summary_writer = SummaryWriter(args.checkpoint_dir) | |
| logger = Logger(lr_scheduler, summary_writer, args.summary_freq, | |
| start_step=start_step) | |
| total_steps = start_step | |
| epoch = start_epoch | |
| print('Start training') | |
| while total_steps < args.num_steps: | |
| model.train() | |
| # mannual change random seed for shuffling every epoch | |
| if args.distributed: | |
| train_sampler.set_epoch(epoch) | |
| for i, sample in enumerate(train_loader): | |
| img1, img2, flow_gt, valid = [x.to(device) for x in sample] | |
| results_dict = model(img1, img2, | |
| attn_splits_list=args.attn_splits_list, | |
| corr_radius_list=args.corr_radius_list, | |
| prop_radius_list=args.prop_radius_list, | |
| ) | |
| flow_preds = results_dict['flow_preds'] | |
| loss, metrics = flow_loss_func(flow_preds, flow_gt, valid, | |
| gamma=args.gamma, | |
| max_flow=args.max_flow, | |
| ) | |
| if isinstance(loss, float): | |
| continue | |
| if torch.isnan(loss): | |
| continue | |
| metrics.update({'total_loss': loss.item()}) | |
| # more efficient zero_grad | |
| for param in model_without_ddp.parameters(): | |
| param.grad = None | |
| loss.backward() | |
| # Gradient clipping | |
| torch.nn.utils.clip_grad_norm_(model.parameters(), args.grad_clip) | |
| optimizer.step() | |
| lr_scheduler.step() | |
| if args.local_rank == 0: | |
| logger.push(metrics) | |
| logger.add_image_summary(img1, img2, flow_preds, flow_gt) | |
| total_steps += 1 | |
| if total_steps % args.save_ckpt_freq == 0 or total_steps == args.num_steps: | |
| if args.local_rank == 0: | |
| checkpoint_path = os.path.join(args.checkpoint_dir, 'step_%06d.pth' % total_steps) | |
| torch.save({ | |
| 'model': model_without_ddp.state_dict() | |
| }, checkpoint_path) | |
| if total_steps % args.save_latest_ckpt_freq == 0: | |
| checkpoint_path = os.path.join(args.checkpoint_dir, 'checkpoint_latest.pth') | |
| if args.local_rank == 0: | |
| torch.save({ | |
| 'model': model_without_ddp.state_dict(), | |
| 'optimizer': optimizer.state_dict(), | |
| 'step': total_steps, | |
| 'epoch': epoch, | |
| }, checkpoint_path) | |
| if total_steps % args.val_freq == 0: | |
| print('Start validation') | |
| val_results = {} | |
| # support validation on multiple datasets | |
| if 'chairs' in args.val_dataset: | |
| results_dict = validate_chairs(model_without_ddp, | |
| with_speed_metric=args.with_speed_metric, | |
| attn_splits_list=args.attn_splits_list, | |
| corr_radius_list=args.corr_radius_list, | |
| prop_radius_list=args.prop_radius_list, | |
| ) | |
| if args.local_rank == 0: | |
| val_results.update(results_dict) | |
| if 'things' in args.val_dataset: | |
| results_dict = validate_things(model_without_ddp, | |
| padding_factor=args.padding_factor, | |
| with_speed_metric=args.with_speed_metric, | |
| attn_splits_list=args.attn_splits_list, | |
| corr_radius_list=args.corr_radius_list, | |
| prop_radius_list=args.prop_radius_list, | |
| ) | |
| if args.local_rank == 0: | |
| val_results.update(results_dict) | |
| if 'sintel' in args.val_dataset: | |
| results_dict = validate_sintel(model_without_ddp, | |
| count_time=args.count_time, | |
| padding_factor=args.padding_factor, | |
| with_speed_metric=args.with_speed_metric, | |
| evaluate_matched_unmatched=args.evaluate_matched_unmatched, | |
| attn_splits_list=args.attn_splits_list, | |
| corr_radius_list=args.corr_radius_list, | |
| prop_radius_list=args.prop_radius_list, | |
| ) | |
| if args.local_rank == 0: | |
| val_results.update(results_dict) | |
| if 'kitti' in args.val_dataset: | |
| results_dict = validate_kitti(model_without_ddp, | |
| padding_factor=args.padding_factor, | |
| with_speed_metric=args.with_speed_metric, | |
| attn_splits_list=args.attn_splits_list, | |
| corr_radius_list=args.corr_radius_list, | |
| prop_radius_list=args.prop_radius_list, | |
| ) | |
| if args.local_rank == 0: | |
| val_results.update(results_dict) | |
| if args.local_rank == 0: | |
| logger.write_dict(val_results) | |
| # Save validation results | |
| val_file = os.path.join(args.checkpoint_dir, 'val_results.txt') | |
| with open(val_file, 'a') as f: | |
| f.write('step: %06d\n' % total_steps) | |
| if args.evaluate_matched_unmatched: | |
| metrics = ['chairs_epe', | |
| 'chairs_s0_10', 'chairs_s10_40', 'chairs_s40+', | |
| 'things_clean_epe', 'things_clean_s0_10', 'things_clean_s10_40', | |
| 'things_clean_s40+', | |
| 'sintel_clean_epe', 'sintel_clean_matched', 'sintel_clean_unmatched', | |
| 'sintel_clean_s0_10', 'sintel_clean_s10_40', | |
| 'sintel_clean_s40+', | |
| 'sintel_final_epe', 'sintel_final_matched', 'sintel_final_unmatched', | |
| 'sintel_final_s0_10', 'sintel_final_s10_40', | |
| 'sintel_final_s40+', | |
| 'kitti_epe', 'kitti_f1', 'kitti_s0_10', 'kitti_s10_40', 'kitti_s40+', | |
| ] | |
| else: | |
| metrics = ['chairs_epe', 'chairs_s0_10', 'chairs_s10_40', 'chairs_s40+', | |
| 'things_clean_epe', 'things_clean_s0_10', 'things_clean_s10_40', | |
| 'things_clean_s40+', | |
| 'sintel_clean_epe', 'sintel_clean_s0_10', 'sintel_clean_s10_40', | |
| 'sintel_clean_s40+', | |
| 'sintel_final_epe', 'sintel_final_s0_10', 'sintel_final_s10_40', | |
| 'sintel_final_s40+', | |
| 'kitti_epe', 'kitti_f1', 'kitti_s0_10', 'kitti_s10_40', 'kitti_s40+', | |
| ] | |
| eval_metrics = [] | |
| for metric in metrics: | |
| if metric in val_results.keys(): | |
| eval_metrics.append(metric) | |
| metrics_values = [val_results[metric] for metric in eval_metrics] | |
| num_metrics = len(eval_metrics) | |
| # save as markdown format | |
| if args.evaluate_matched_unmatched: | |
| f.write(("| {:>25} " * num_metrics + '\n').format(*eval_metrics)) | |
| f.write(("| {:25.3f} " * num_metrics).format(*metrics_values)) | |
| else: | |
| f.write(("| {:>20} " * num_metrics + '\n').format(*eval_metrics)) | |
| f.write(("| {:20.3f} " * num_metrics).format(*metrics_values)) | |
| f.write('\n\n') | |
| model.train() | |
| if total_steps >= args.num_steps: | |
| print('Training done') | |
| return | |
| epoch += 1 | |
| if __name__ == '__main__': | |
| parser = get_args_parser() | |
| args = parser.parse_args() | |
| if 'LOCAL_RANK' not in os.environ: | |
| os.environ['LOCAL_RANK'] = str(args.local_rank) | |
| main(args) | |