def validate_sintel(model, iters=32):
""" Peform validation using the Sintel (train) split """
model.eval()
results = {}
for dstype in ['clean', 'final']:
val_dataset = datasets.MpiSintel(split='training', dstype=dstype)
epe_list = []
for val_id in range(len(val_dataset)):
image1, image2, flow_gt, _ = val_dataset[val_id]
image1 = image1[None].cuda()
image2 = image2[None].cuda()
padder = InputPadder(image1.shape)
image1, image2 = padder.pad(image1, image2)
flow_low, flow_pr = model(image1, image2, iters=iters, test_mode=True)
flow = padder.unpad(flow_pr[0]).cpu()
epe = torch.sum((flow - flow_gt)**2, dim=0).sqrt()
epe_list.append(epe.view(-1).numpy())
epe_all = np.concatenate(epe_list)
epe = np.mean(epe_all)
px1 = np.mean(epe_all<1)
px3 = np.mean(epe_all<3)
px5 = np.mean(epe_all<5)
print("Validation (%s) EPE: %f, 1px: %f, 3px: %f, 5px: %f" % (dstype, epe, px1, px3, px5))
results[dstype] = np.mean(epe_list)
return results
def create_sintel_submission(model, iters=32, warm_start=False, output_path='sintel_submission'):
""" Create submission for the Sintel leaderboard """
model.eval()
for dstype in ['clean', 'final']:
test_dataset = datasets.MpiSintel(split='test', aug_params=None, dstype=dstype)
flow_prev, sequence_prev = None, None
for test_id in range(len(test_dataset)):
image1, image2, (sequence, frame) = test_dataset[test_id]
if sequence != sequence_prev:
flow_prev = None
padder = InputPadder(image1.shape)
image1, image2 = padder.pad(image1[None].cuda(), image2[None].cuda())
flow_low, flow_pr = model(image1, image2, iters=iters, flow_init=flow_prev, test_mode=True)
flow = padder.unpad(flow_pr[0]).permute(1, 2, 0).cpu().numpy()
if warm_start:
flow_prev = forward_interpolate(flow_low[0])[None].cuda()
output_dir = os.path.join(output_path, dstype, sequence)
output_file = os.path.join(output_dir, 'frame%04d.flo' % (frame+1))
if not os.path.exists(output_dir):
os.makedirs(output_dir)
frame_utils.writeFlow(output_file, flow)
sequence_prev = sequence
def validate_sintel(args, model, iters=50):
""" Evaluate trained model on Sintel(train) clean + final passes """
model.eval()
pad = 2
for dstype in ['clean', 'final']:
val_dataset = datasets.MpiSintel(args,
do_augument=False,
dstype=dstype)
epe_list = []
for i in range(len(val_dataset)):
image1, image2, flow_gt, _ = val_dataset[i]
image1 = image1[None].cuda()
image2 = image2[None].cuda()
image1 = F.pad(image1, [0, 0, pad, pad], mode='replicate')
image2 = F.pad(image2, [0, 0, pad, pad], mode='replicate')
with torch.no_grad():
flow_predictions = model.module(image1, image2, iters=iters)
flow_pr = flow_predictions[-1][0, :, pad:-pad]
epe = torch.sum((flow_pr - flow_gt.cuda())**2, dim=0)
epe = torch.sqrt(epe).mean()
epe_list.append(epe.item())
print("Validation (%s) EPE: %f" % (dstype, np.mean(epe_list)))
def fetch_dataloader(args):
""" Create the data loader for the corresponding trainign set """
if args.dataset == 'chairs':
train_dataset = datasets.FlyingChairs(args, image_size=args.image_size)
elif args.dataset == 'things':
clean_dataset = datasets.SceneFlow(args,
image_size=args.image_size,
dstype='frames_cleanpass')
final_dataset = datasets.SceneFlow(args,
image_size=args.image_size,
dstype='frames_finalpass')
train_dataset = clean_dataset + final_dataset
elif args.dataset == 'sintel':
clean_dataset = datasets.MpiSintel(args,
image_size=args.image_size,
dstype='clean')
final_dataset = datasets.MpiSintel(args,
image_size=args.image_size,
dstype='final')
train_dataset = clean_dataset + final_dataset
elif args.dataset == 'kitti':
train_dataset = datasets.KITTI(args,
image_size=args.image_size,
is_val=False)
gpuargs = {'num_workers': 4, 'drop_last': True}
train_loader = DataLoader(train_dataset,
batch_size=args.batch_size,
pin_memory=True,
shuffle=True,
**gpuargs)
print('Training with %d image pairs' % len(train_dataset))
return train_loader
def main():
global args, best_EPE, save_path
args = parser.parse_args()
# Load config file
if args.cfg is not None:
cfg_from_file(args.cfg)
assert cfg.TAG == os.path.splitext(os.path.basename(
args.cfg))[0], 'TAG name should be file name'
# Build save_path, which can be specified by out_dir and exp_dir
save_path = '{},{}epochs{},b{},lr{}'.format(
'dicl_wrapper', args.epochs,
',epochSize' + str(args.epoch_size) if args.epoch_size > 0 else '',
args.batch_size, args.lr)
save_path = os.path.join(args.exp_dir, save_path)
if args.out_dir is not None:
outpath = os.path.join(args.out_dir, args.dataset)
else:
outpath = args.dataset
save_path = os.path.join(outpath, save_path)
if not os.path.exists(outpath): os.makedirs(outpath)
if not os.path.exists(save_path): os.makedirs(save_path)
# Create logger
log_file = os.path.join(save_path, 'log.txt')
logger = create_logger(log_file)
logger.info('**********************Start logging**********************')
logger.info('=> will save everything to {}'.format(save_path))
# Print settings
for _, key in enumerate(args.__dict__):
logger.info(args.__dict__[key])
save_config_to_file(cfg, logger=logger)
logger.info(args.pretrained)
# Set random seed
torch.cuda.manual_seed(args.seed)
np.random.seed(args.seed)
train_writer = SummaryWriter(os.path.join(save_path, 'train'))
eval_writer = SummaryWriter(os.path.join(save_path, 'eval'))
logger.info("=> fetching img pairs in '{}'".format(args.data))
########################## DATALOADER ##########################
if args.dataset == 'flying_chairs':
if cfg.SIMPLE_AUG:
train_dataset = datasets.FlyingChairs_SimpleAug(args,
root=args.data)
test_dataset = datasets.FlyingChairs_SimpleAug(args,
root=args.data,
mode='val')
else:
train_dataset = datasets.FlyingChairs(args,
image_size=cfg.CROP_SIZE,
root=args.data)
test_dataset = datasets.FlyingChairs(args,
root=args.data,
mode='val',
do_augument=False)
elif args.dataset == 'flying_things':
train_dataset = datasets.SceneFlow(args,
image_size=cfg.CROP_SIZE,
root=args.data,
dstype='frames_cleanpass',
mode='train')
test_dataset = datasets.SceneFlow(args,
image_size=cfg.CROP_SIZE,
root=args.data,
dstype='frames_cleanpass',
mode='val',
do_augument=False)
elif args.dataset == 'mpi_sintel_clean' or args.dataset == 'mpi_sintel_final':
clean_dataset = datasets.MpiSintel(args,
image_size=cfg.CROP_SIZE,
root=args.data,
dstype='clean')
final_dataset = datasets.MpiSintel(args,
image_size=cfg.CROP_SIZE,
root=args.data,
dstype='final')
train_dataset = torch.utils.data.ConcatDataset([clean_dataset] +
[final_dataset])
if args.dataset == 'mpi_sintel_final':
test_dataset = datasets.MpiSintel(args,
do_augument=False,
image_size=None,
root=args.data,
dstype='final')
else:
test_dataset = datasets.MpiSintel(args,
do_augument=False,
image_size=None,
root=args.data,
dstype='clean')
elif args.dataset == 'KITTI':
train_dataset = datasets.KITTI(args,
image_size=cfg.CROP_SIZE,
root=args.data,
is_val=False,
logger=logger)
if args.data_kitti12 is not None:
train_dataset12 = datasets.KITTI12(args,
image_size=cfg.CROP_SIZE,
root=args.data_kitti12,
is_val=False,
logger=logger)
train_dataset = torch.utils.data.ConcatDataset([train_dataset] +
[train_dataset12])
test_dataset = datasets.KITTI(args,
root=args.data,
do_augument=False,
is_val=True,
do_pad=False)
else:
raise NotImplementedError
logger.info('Training with %d image pairs' % len(train_dataset))
logger.info('Testing with %d image pairs' % len(test_dataset))
gpuargs = {'num_workers': args.workers, 'drop_last': cfg.DROP_LAST}
train_loader = torch.utils.data.DataLoader(train_dataset,
batch_size=args.batch_size,
pin_memory=True,
shuffle=True,
**gpuargs)
if 'KITTI' in args.dataset:
# We set batch size to 1 since KITTI images have different sizes
val_loader = torch.utils.data.DataLoader(test_dataset,
batch_size=1,
num_workers=args.workers,
pin_memory=True,
shuffle=False)
else:
val_loader = torch.utils.data.DataLoader(test_dataset,
batch_size=args.batch_size,
num_workers=args.workers,
pin_memory=True,
shuffle=False)
# create model
if args.pretrained:
logger.info("=> using pre-trained model '{}'".format(args.pretrained))
pretrained_dict = torch.load(args.pretrained)
if 'state_dict' in pretrained_dict.keys():
pretrained_dict['state_dict'] = {
k: v
for k, v in pretrained_dict['state_dict'].items()
}
model = models.__dict__['dicl_wrapper'](None)
assert (args.solver in ['adam', 'sgd'])
logger.info('=> setting {} solver'.format(args.solver))
if args.solver == 'adam':
optimizer = torch.optim.Adam(model.parameters(),
lr=args.lr,
weight_decay=cfg.WEIGHT_DECAY,
betas=(cfg.MOMENTUM, cfg.BETA))
elif args.solver == 'sgd':
optimizer = torch.optim.SGD(model.parameters(),
lr=args.lr,
weight_decay=cfg.WEIGHT_DECAY,
momentum=cfg.MOMENTUM)
if args.pretrained:
if 'state_dict' in pretrained_dict.keys():
model.load_state_dict(pretrained_dict['state_dict'], strict=False)
else:
model.load_state_dict(pretrained_dict, strict=False)
if args.reuse_optim:
try:
optimizer.load_state_dict(pretrained_dict['optimizer_state'])
except:
logger.info('do not have optimizer state')
del pretrained_dict
torch.cuda.empty_cache()
model = torch.nn.DataParallel(model)
if torch.cuda.is_available():
model = model.cuda()
# Evaluation
if args.evaluate:
with torch.no_grad():
best_EPE = validate(val_loader,
model,
0,
None,
eval_writer,
logger=logger)
return
# Learning rate schedule
milestones = []
for num in range(len(args.milestones)):
milestones.append(int(args.milestones[num]))
scheduler = torch.optim.lr_scheduler.MultiStepLR(optimizer,
milestones=milestones,
gamma=0.5)
###################################### Training ######################################
for epoch in range(args.start_epoch, args.epochs):
# train for one epoch
train_loss = train(train_loader,
model,
optimizer,
epoch,
train_writer,
logger=logger)
scheduler.step()
train_writer.add_scalar('lr', optimizer.param_groups[0]['lr'], epoch)
train_writer.add_scalar('avg_loss', train_loss, epoch)
if epoch % args.eval_freq == 0 and not args.no_eval:
with torch.no_grad():
EPE = validate(val_loader,
model,
epoch,
output_writers,
eval_writer,
logger=logger)
eval_writer.add_scalar('mean_EPE', EPE, epoch)
if best_EPE < 0:
best_EPE = EPE
if EPE < best_EPE:
best_EPE = EPE
ckpt_best_file = 'checkpoint_best.pth.tar'
save_checkpoint(
{
'epoch': epoch + 1,
'arch': 'dicl_wrapper',
'state_dict': model.module.state_dict(),
'optimizer_state': optimizer.state_dict(),
'best_EPE': EPE
},
False,
filename=ckpt_best_file)
logger.info('Epoch: [{0}] Best EPE: {1}'.format(epoch, best_EPE))
# Skip at least 5 epochs to save memory
save_freq = max(args.eval_freq, 5)
if epoch % save_freq == 0:
ckpt_file = 'checkpoint_' + str(epoch) + '.pth.tar'
save_checkpoint(
{
'epoch': epoch + 1,
'arch': 'dicl_wrapper',
'state_dict': model.module.state_dict(),
'optimizer_state': optimizer.state_dict(),
'best_EPE': best_EPE
},
False,
filename=ckpt_file)
