def fetch_dataloader(args):
""" Create the data loader for the corresponding trainign set """
if args.dataset == 'chairs':
train_dataset = datasets.FlyingChairs(args,
root=args.dataset_root,
image_size=args.image_size)
elif args.dataset == 'things':
clean_dataset = datasets.SceneFlow(args,
root=args.dataset_root,
image_size=args.image_size,
dstype='frames_cleanpass')
final_dataset = datasets.SceneFlow(args,
root=args.dataset_root,
image_size=args.image_size,
dstype='frames_finalpass')
train_dataset = clean_dataset + final_dataset
elif args.dataset == 'sintel':
clean_dataset = datasets.MpiSintel_Train(args,
root=args.dataset_root,
image_size=args.image_size,
dstype='clean')
final_dataset = datasets.MpiSintel_Train(args,
root=args.dataset_root,
image_size=args.image_size,
dstype='final')
assert len(clean_dataset) == 908 and len(final_dataset) == 908
train_dataset = clean_dataset + final_dataset
elif args.dataset == 'kitti':
train_dataset = datasets.KITTI(args,
root=args.dataset_root,
image_size=args.image_size,
is_val=False)
else:
raise NotImplementedError
gpuargs = {'num_workers': args.num_of_workers, '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 validate_chairs(model, iters=24):
""" Perform evaluation on the FlyingChairs (test) split """
model.eval()
epe_list = []
val_dataset = datasets.FlyingChairs(split='validation')
for val_id in range(len(val_dataset)):
image1, image2, flow_gt, _ = val_dataset[val_id]
image1 = image1[None].cuda()
image2 = image2[None].cuda()
_, flow_pr = model(image1, image2, iters=iters, test_mode=True)
epe = torch.sum((flow_pr[0].cpu() - flow_gt)**2, dim=0).sqrt()
epe_list.append(epe.view(-1).numpy())
epe = np.mean(np.concatenate(epe_list))
print("Validation Chairs EPE: %f" % epe)
return {'chairs': epe}
def main():
global args, best_EPE
args = parser.parse_args()
# create model
if args.pretrained:
print("=> using pre-trained model '{}'".format(args.arch))
else:
print("=> creating model '{}'".format(args.arch))
model = models.__dict__[args.arch](args.pretrained)
model = torch.nn.DataParallel(model).cuda()
criterion = multiscaleloss().cuda()
high_res_EPE = multiscaleloss(scales=1,
downscale=4,
weights=(1),
loss='L1').cuda()
cudnn.benchmark = True
# Data loading code
normalize = transforms.Normalize(mean=[0.485, 0.456, 0.406],
std=[0.229, 0.224, 0.225])
dataset = datasets.FlyingChairs(
args.data,
transform=transforms.Compose([transforms.ToTensor(), normalize]),
target_transform=None,
co_transform=flow_transforms.Compose([
flow_transforms.RandomTranslate(10),
flow_transforms.RandomCropRotate(10, 360, 5),
flow_transforms.RandomCrop((320, 448))
]),
split=args.split)
train_loader = torch.utils.data.DataLoader(dataset,
batch_size=args.batch_size,
shuffle=True,
num_workers=args.workers,
pin_memory=True)
dataset.eval()
val_loader = torch.utils.data.DataLoader(dataset,
batch_size=args.batch_size,
shuffle=False,
num_workers=args.workers,
pin_memory=True)
optimizer = torch.optim.SGD(model.parameters(),
args.lr,
momentum=args.momentum,
weight_decay=args.weight_decay)
if args.evaluate:
best_EPE = validate(val_loader, model, criterion, high_res_EPE)
return
for epoch in range(args.start_epoch, args.epochs):
adjust_learning_rate(optimizer, epoch)
# train for one epoch
train(train_loader, model, criterion, high_res_EPE, optimizer, epoch)
# evaluate o validation set
EPE = validate(val_loader, model, criterion, high_res_EPE)
if best_EPE < 0:
best_EPE = EPE
# remember best prec@1 and save checkpoint
is_best = EPE < best_EPE
best_EPE = min(EPE, best_EPE)
save_checkpoint(
{
'epoch': epoch + 1,
'arch': args.arch,
'state_dict': model.state_dict(),
'best_EPE': best_EPE,
}, is_best)
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)
