def train(self, model, loader, optimizer, criterion, lr_init=1e-2, epochs=3000,
swag_model=None, swag=False, swag_start=2000, swag_freq=50, swag_lr=1e-3,
print_freq=100, use_cuda=False, const_lr=False):
# copied from pavels regression notebook
if const_lr:
lr = lr_init
train_res_list = []
for epoch in range(epochs):
if not const_lr:
t = (epoch + 1) / swag_start if swag else (epoch + 1) / epochs
lr_ratio = swag_lr / lr_init if swag else 0.05
if t <= 0.5:
factor = 1.0
elif t <= 0.9:
factor = 1.0 - (1.0 - lr_ratio) * (t - 0.5) / 0.4
else:
factor = lr_ratio
lr = factor * lr_init
adjust_learning_rate(optimizer, factor)
train_res = utils.train_epoch(loader, model, criterion, optimizer, cuda=use_cuda, regression=True)
train_res_list.append(train_res)
if swag and epoch > swag_start:
swag_model.collect_model(model)
if (epoch % print_freq == 0 or epoch == epochs - 1):
print('Epoch %d. LR: %g. Loss: %.4f' % (epoch, lr, train_res['loss']))
return train_res_list
def fit(self, *args, **kwargs):
for epoch in range(self.epochs):
train_res = train_epoch(self.loader,
self,
self.elbo,
self.optimizer,
regression=True,
cuda=self.use_cuda)
values = [
'%d/%d' % (epoch + 1, self.epochs), train_res['accuracy'],
train_res['loss'], train_res['stats']['kl'],
train_res['stats']['nll']
]
def fit(self, mean, variance, cov_factor, loader, criterion, epochs=100):
print('Fitting VI')
self.vi_model.set_subspace(mean, cov_factor)
elbo = ELBO(criterion, len(loader.dataset))
optimizer = torch.optim.Adam([
param for param in self.vi_model.parameters()
if param.requires_grad
])
for _ in range(epochs):
train_res = train_epoch(loader, self.vi_model, elbo, optimizer)
print(train_res)
def fit(self, mean, variance, subspace, use_cuda=True, **kwargs):
if use_cuda and torch.cuda.is_available():
self.mean = mean.cuda()
self.subspace = subspace.cuda()
else:
self.mean = mean
self.subspace = subspace
if self.proj_params is None:
proj_params = torch.zeros(self.subspace.size(0),
1,
dtype=self.subspace.dtype,
device=self.subspace.device,
requires_grad=True)
print(proj_params.device)
self.proj_model = ProjectedModel(model=self.model,
mean=self.mean.unsqueeze(1),
projection=self.subspace,
proj_params=proj_params)
# define optimizer
self.optimizer = torch.optim.SGD([proj_params], **self.kwargs)
else:
proj_params = self.proj_params.clone()
# now train projected parameters
loss_vec = []
for _ in range(self.epochs):
loss = train_epoch(loader=self.loader,
optimizer=self.optimizer,
model=self.proj_model,
criterion=self.criterion,
**kwargs)
loss_vec.append(loss)
self.proj_params = proj_params
return loss_vec
# optimizer = torch.optim.Adam(vi_model.parameters(),
# lr=args.learning_rate,
# weight_decay=0.0005)
# else:
# optimizer = torch.optim.SGD(vi_model.parameters(),
# lr=args.learning_rate,
# momentum=args.momentum,
# weight_decay=args.weight_decay)
#printf, logfile = utils.get_logging_print(os.path.join(args.dir, args.log_fname + '-%s.txt'))
#print('Saving logs to: %s' % logfile)
columns = ['ep', 'acc', 'loss', 'kl', 'nll', 'sigma_1', 'time']
epoch = 0
for epoch in range(args.epochs):
time_ep = time.time()
train_res = utils.train_epoch(train_loader, vi_model, elbo, optimizer,
epoch)
time_ep = time.time() - time_ep
sigma_1 = torch.nn.functional.softplus(
vi_model.inv_softplus_sigma.detach().cpu())[0].item()
values = [
'%d/%d' % (epoch + 1, args.epochs), train_res['accuracy'],
train_res['loss'], train_res['stats']['kl'], train_res['stats']['nll'],
sigma_1, time_ep
]
if epoch == 0:
print(
tabulate.tabulate([values],
columns,
tablefmt='simple',
floatfmt='8.4f'))
else:
sgd_ens_preds = None
sgd_targets = None
n_ensembled = 0.
for epoch in range(start_epoch, args.epochs):
time_ep = time.time()
if not args.no_schedule:
lr = schedule(epoch, args)
utils.adjust_learning_rate(optimizer, lr)
else:
lr = args.lr_init
if (args.swag and (epoch + 1) > args.swag_start) and args.cov_mat:
train_res = utils.train_epoch(train_loader, model, criterion, optimizer, epoch)
else:
train_res = utils.train_epoch(train_loader, model, criterion, optimizer, epoch)
if epoch == 0 or epoch % args.eval_freq == args.eval_freq - 1 or epoch == args.epochs - 1:
test_res = utils.eval(val_loader, model, criterion)
utils.mata_eval(model, meta_valloader, meta_testloader, 'BASELINE')
else:
test_res = {'loss': None, 'accuracy': None}
if args.swag and (epoch + 1) > args.swag_start and (epoch + 1 - args.swag_start) % args.swag_c_epochs == 0:
sgd_preds, sgd_targets = utils.predictions(val_loader, model)
# sgd_res = utils.predict(val_loader, model)
# sgd_preds = sgd_res["predictions"]
# sgd_targets = sgd_res["targets"]
elbo = ELBO(losses.cross_entropy, len(loaders['train'].dataset),
args.temperature)
#optimizer = torch.optim.Adam([param for param in vi_model.parameters()], lr=0.01)
optimizer = torch.optim.SGD([param for param in vi_model.parameters()],
lr=args.lr,
momentum=0.9)
#printf, logfile = utils.get_logging_print(os.path.join(args.dir, args.log_fname + '-%s.txt'))
#print('Saving logs to: %s' % logfile)
columns = ['ep', 'acc', 'loss', 'kl', 'nll', 'sigma_1', 'time']
epoch = 0
for epoch in range(args.epochs):
time_ep = time.time()
train_res = utils.train_epoch(loaders['train'], vi_model, elbo, optimizer)
time_ep = time.time() - time_ep
sigma_1 = torch.nn.functional.softplus(
vi_model.inv_softplus_sigma.detach().cpu())[0].item()
values = [
'%d/%d' % (epoch + 1, args.epochs), train_res['accuracy'],
train_res['loss'], train_res['stats']['kl'], train_res['stats']['nll'],
sigma_1, time_ep
]
if epoch == 0:
print(
tabulate.tabulate([values],
columns,
tablefmt='simple',
floatfmt='8.4f'))
else:
sgd_ens_preds = None
sgd_targets = None
n_ensembled = 0.
for epoch in range(start_epoch, args.epochs):
time_ep = time.time()
if not args.no_schedule:
lr = schedule(epoch)
utils.adjust_learning_rate(optimizer, lr)
else:
lr = args.lr_init
if (args.swag and (epoch + 1) > args.swag_start) and args.cov_mat:
train_res = utils.train_epoch(loaders['train'], model, criterion,
optimizer)
else:
train_res = utils.train_epoch(loaders['train'], model, criterion,
optimizer)
if epoch == 0 or epoch % args.eval_freq == args.eval_freq - 1 or epoch == args.epochs - 1:
test_res = utils.eval(loaders['test'], model, criterion)
else:
test_res = {'loss': None, 'accuracy': None}
if args.swag and (epoch + 1) > args.swag_start and (
epoch + 1 - args.swag_start) % args.swag_c_epochs == 0:
#sgd_preds, sgd_targets = utils.predictions(loaders["test"], model)
sgd_res = utils.predict(loaders["test"], model)
sgd_preds = sgd_res["predictions"]
sgd_targets = sgd_res["targets"]