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)