# Optimization
def tensor_in(t, a):
for a_ in a:
if t is a_:
return True
return False
scheduler = None
if args.optimizer == 'adam':
optimizer = optim.Adam(model.parameters(), lr=args.lr, betas=(0.9, 0.99), weight_decay=args.wd)
if args.scheduler: scheduler = CosineAnnealingWarmRestarts(optimizer, 20, T_mult=2, last_epoch=args.begin_epoch - 1)
elif args.optimizer == 'adamax':
optimizer = optim.Adamax(model.parameters(), lr=args.lr, betas=(0.9, 0.99), weight_decay=args.wd)
elif args.optimizer == 'rmsprop':
optimizer = optim.RMSprop(model.parameters(), lr=args.lr, weight_decay=args.wd)
elif args.optimizer == 'sgd':
optimizer = torch.optim.SGD(model.parameters(), lr=args.lr, momentum=0.9, weight_decay=args.wd)
if args.scheduler:
scheduler = torch.optim.lr_scheduler.MultiStepLR(
optimizer, milestones=[60, 120, 160], gamma=0.2, last_epoch=args.begin_epoch - 1
)
else:
raise ValueError('Unknown optimizer {}'.format(args.optimizer))
best_test_bpd = math.inf
if (args.resume is not None):
logger.info('Resuming model from {}'.format(args.resume))
with torch.no_grad():
# params = [par for par in gmm.parameters()]
if args.optimizer == 'adam':
optimizer = optim.Adam(params,
lr=args.lr,
betas=(0.9, 0.99),
weight_decay=args.wd)
if args.scheduler:
scheduler = CosineAnnealingWarmRestarts(optimizer,
20,
T_mult=2,
last_epoch=args.begin_epoch -
1)
elif args.optimizer == 'adamax':
optimizer = optim.Adamax(params,
lr=args.lr,
betas=(0.9, 0.99),
weight_decay=args.wd)
elif args.optimizer == 'rmsprop':
optimizer = optim.RMSprop(params, lr=args.lr, weight_decay=args.wd)
elif args.optimizer == 'sgd':
optimizer = torch.optim.SGD(params,
lr=args.lr,
momentum=0.9,
weight_decay=args.wd)
if args.scheduler:
scheduler = torch.optim.lr_scheduler.MultiStepLR(
optimizer,
milestones=[60, 120, 160],
gamma=0.2,
last_epoch=args.begin_epoch - 1)
else:
