model = MapNet(mapnet=posenet)
else:
raise NotImplementedError
# loss function
if args.model == 'posenet':
train_criterion = PoseNetCriterion(sax=sax, saq=saq, learn_beta=True)
val_criterion = PoseNetCriterion()
elif args.model.find('mapnet') >= 0:
kwargs = dict(sax=sax,
saq=saq,
srx=srx,
srq=srq,
learn_beta=True,
learn_gamma=True)
train_criterion = MapNetCriterion(**kwargs)
val_criterion = MapNetCriterion()
else:
raise NotImplementedError
# optimizer
param_list = [{'params': model.parameters()}]
if hasattr(train_criterion, 'sax') and hasattr(train_criterion, 'saq'):
param_list.append({'params': [train_criterion.sax, train_criterion.saq]})
if hasattr(train_criterion, 'srx') and hasattr(train_criterion, 'srq'):
param_list.append({'params': [train_criterion.srx, train_criterion.srq]})
optimizer = Optimizer(params=param_list,
method=opt_method,
base_lr=lr,
weight_decay=weight_decay,
**optim_config)
dual_target=True,
sas=sas,
learn_sigma=args.learn_sigma)
if '++' in args.model:
kwargs = dict(kwargs, gps_mode=(vo_lib == 'gps'))
train_criterion = MapNetOnlineCriterion(**kwargs)
val_criterion = MapNetOnlineCriterion()
elif args.uncertainty_criterion:
train_criterion = UncertainyCriterion(
**kwargs, learn_log=not args.learn_direct_sigma)
val_criterion = UncertainyCriterion(
dual_target='multitask' in args.model,
learn_log=not args.learn_direct_sigma)
else:
train_criterion = MapNetCriterion(**kwargs)
val_criterion = MapNetCriterion(dual_target='multitask' in args.model)
else:
raise NotImplementedError
# optimizer
param_list = [{'params': model.parameters()}]
if args.learn_beta and hasattr(train_criterion, 'sax') and \
hasattr(train_criterion, 'saq'):
param_list.append({'params': [train_criterion.sax, train_criterion.saq]})
if args.learn_gamma and hasattr(train_criterion, 'srx') and \
hasattr(train_criterion, 'srq'):
param_list.append({'params': [train_criterion.srx, train_criterion.srq]})
if args.learn_sigma and hasattr(train_criterion, 'sas'):
param_list.append({'params': [train_criterion.sas]})
optimizer = Optimizer(params=param_list,