def train_epoch(model,
loaders,
criterion,
optimizer,
epoch,
end_epoch,
eval_freq=1,
save_freq=10,
output_dir='./',
lr_init=0.01):
time_ep = time.time()
lr = training_utils.schedule(epoch, lr_init, end_epoch, swa=False)
training_utils.adjust_learning_rate(optimizer, lr)
train_res = training_utils.train_epoch(loaders["train"], model, criterion,
optimizer)
if (epoch == 0 or epoch % eval_freq == eval_freq - 1
or epoch == end_epoch - 1):
test_res = training_utils.eval(loaders["test"], model, criterion)
else:
test_res = {"loss": None, "accuracy": None}
if (epoch + 1) % save_freq == 0:
training_utils.save_checkpoint(
output_dir,
epoch + 1,
state_dict=model.state_dict(),
optimizer=optimizer.state_dict(),
)
time_ep = time.time() - time_ep
values = [
epoch + 1,
lr,
train_res["loss"],
train_res["accuracy"],
test_res["loss"],
test_res["accuracy"],
time_ep,
]
table = tabulate.tabulate([values],
columns,
tablefmt="simple",
floatfmt="8.4f")
if epoch % 40 == 0:
table = table.split("\n")
table = "\n".join([table[1]] + table)
else:
table = table.split("\n")[2]
print(table)
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, 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
lr=5e-4,
momentum=0.9,
weight_decay=0)
swag_model.sample(0)
utils.bn_update(loaders['train'], swag_model)
print(utils.eval(loaders['test'], swag_model, criterion))
printf, logfile = utils.get_logging_print(
os.path.join(args.dir, args.log_fname + '-%s.txt'))
print('Saving logs to: %s' % logfile)
#printf=print
columns = ['ep', 'acc', 'loss', 'prior']
for epoch in range(args.epochs):
train_res = utils.train_epoch(loaders['train'], proj_model, criterion,
optimizer)
values = [
'%d/%d' % (epoch + 1, args.epochs), train_res['accuracy'],
train_res['loss'], train_res['stats']['prior'],
train_res['stats']['nll']
]
if epoch == 0:
printf(
tabulate.tabulate([values],
columns,
tablefmt='simple',
floatfmt='8.4f'))
else:
printf(
tabulate.tabulate([values],
columns,
vi_model = vi_model.cuda()
print(utils.eval(loaders["train"], vi_model, criterion=losses.cross_entropy))
elbo = ELBO_NF(losses.cross_entropy_output, len(loaders["train"].dataset),
args.temperature)
optimizer = torch.optim.Adam([param for param in vi_model.parameters()],
lr=0.01)
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']
for epoch in range(args.epochs):
train_res = utils.train_epoch(loaders['train'], vi_model, elbo, optimizer)
values = [
'%d/%d' % (epoch + 1, args.epochs), train_res['accuracy'],
train_res['loss'], train_res['stats']['kl'], train_res['stats']['nll']
]
if epoch == 0:
printf(
tabulate.tabulate([values],
columns,
tablefmt='simple',
floatfmt='8.4f'))
else:
printf(
tabulate.tabulate([values],
columns,
tablefmt='plain',
sgd_ens_preds = None
sgd_targets = None
n_ensembled = 0.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.swa and (epoch + 1) > args.swa_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.swa and (epoch + 1) > args.swa_start
and (epoch + 1 - args.swa_start) % args.swa_c_epochs == 0):
# sgd_preds, sgd_targets = utils.predictions(loaders["test"], model)
sgd_res = utils.predict(loaders["test"], model)
sgd_preds = sgd_res["predictions"]
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
print('EPOCH %d. TRAIN' % (epoch + 1))
if args.swa and (epoch + 1) > args.swa_start:
subset = 1.0 / args.swa_freq
for i in range(args.swa_freq):
print('PART %d/%d' % (i + 1, args.swa_freq))
train_res = utils.train_epoch(loaders['train'], model, criterion, optimizer, subset=subset, verbose=True)
num_iterates += 1
utils.save_checkpoint(
args.dir,
num_iterates,
name='iter',
state_dict=model.state_dict(),
)
model.to(args.swa_device)
swag_model.collect_model(model)
model.to(args.device)
else:
train_res = utils.train_epoch(loaders['train'], model, criterion, optimizer, verbose=True)
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, regularizer=regularizer)
else:
train_res = utils.train_epoch(loaders['train'], model, criterion, optimizer, regularizer=regularizer)
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"]
# print("updating sgd_ens")
if sgd_ens_preds is None:
sgd_ens_preds = None
sgd_targets = None
n_ensembled = 0.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.swa and (epoch + 1) > args.swa_start) and args.cov_mat:
train_res = utils.train_epoch(loaders["train"], model, criterion, optimizer, cuda=use_cuda)
else:
train_res = utils.train_epoch(loaders["train"], model, criterion, optimizer, cuda=use_cuda)
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, cuda=use_cuda)
else:
test_res = {"loss": None, "accuracy": None}
if (
args.swa
and (epoch + 1) > args.swa_start
