targets = np.zeros((query_length, args.n_ways))
#printf, logfile = utils.get_logging_print(os.path.join(args.dir, args.log_fname + '-%s.txt'))
#print('Saving logs to: %s' % logfile)
columns = ['iter ens', 'acc', 'nll']
for i in range(num_samples):
with torch.no_grad():
w = vi_model.sample()
offset = 0
for param in eval_model.parameters():
param.data.copy_(w[offset:offset + param.numel()].view(
param.size()).to(args.device))
offset += param.numel()
utils.bn_update(train_loader, eval_model, subset=args.bn_subset)
pred_res = utils.predict(meta_testloader, eval_model)
utils.mata_eval(eval_model, meta_valloader, meta_testloader, 'SWAG-vi')
ens_predictions += pred_res['predictions']
targets = pred_res['targets']
values = [
'%3d/%3d' % (i + 1, num_samples),
np.mean(np.argmax(ens_predictions, axis=1) == targets),
nll(ens_predictions / (i + 1), targets)
]
table = tabulate.tabulate([values],
columns,
tablefmt='simple',
floatfmt='8.4f')
cov_factor = np.vstack((u[None, :], v[None, :]))
subspace = SubspaceModel(torch.FloatTensor(mean),
torch.FloatTensor(cov_factor))
coords = np.dot(
cov_factor / np.sum(np.square(cov_factor), axis=1, keepdims=True),
(w - mean[None, :]).T).T
theta = torch.FloatTensor(coords[2, :])
for i in range(3):
v = subspace(torch.FloatTensor(coords[i]))
offset = 0
for param in model.parameters():
param.data.copy_(v[offset:offset + param.numel()].view(
param.size()).to(args.device))
offset += param.numel()
utils.bn_update(train_loader, model)
print("Performance of model", i, "on the curve", end=":")
utils.mata_eval(model, meta_valloader, meta_testloader, 'curve')
else:
assert len(args.checkpoint) == 1
if args.subspace == 'MNPCA_MANY':
swag_model = SWAG(
model_cfg.base,
args.subspace,
{
'max_rank': 20,
'cov_mat': args.sample_collect
#'pca_rank': args.rank,
},
1e-6,
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"]
# print("updating sgd_ens")
if sgd_ens_preds is None:
sgd_ens_preds = sgd_preds.copy()
else:
#TODO: rewrite in a numerically stable way
sgd_ens_preds = sgd_ens_preds * n_ensembled / (n_ensembled + 1) + sgd_preds/ (n_ensembled + 1)
n_ensembled += 1
swag_model.collect_model(model)
if epoch == 0 or epoch % args.eval_freq == args.eval_freq - 1 or epoch == args.epochs - 1:
# swag_model.set_swa()
swag_model.sample(0.0)
utils.bn_update(train_loader, swag_model)
swag_res = utils.eval(val_loader, swag_model, criterion)
utils.mata_eval(swag_model, meta_valloader, meta_testloader, 'SWAG', classifier='LR')
else:
swag_res = {'loss': None, 'accuracy': None}
if (epoch + 1) % args.save_freq == 0:
utils.save_checkpoint(
args.dir,
epoch + 1,
state_dict=model.state_dict(),
optimizer=optimizer.state_dict()
)
if args.swag:
utils.save_checkpoint(
args.dir,
targets = np.zeros(len(loaders['test'].dataset))
#printf, logfile = utils.get_logging_print(os.path.join(args.dir, args.log_fname + '-%s.txt'))
#print('Saving logs to: %s' % logfile)
columns = ['iter ens', 'acc', 'nll']
for i in range(num_samples):
with torch.no_grad():
w = vi_model.sample()
offset = 0
for param in eval_model.parameters():
param.data.copy_(w[offset:offset + param.numel()].view(
param.size()).to(args.device))
offset += param.numel()
utils.bn_update(loaders['train'], eval_model, subset=args.bn_subset)
pred_res = utils.predict(loaders['test'], eval_model)
ens_predictions += pred_res['predictions']
targets = pred_res['targets']
values = [
'%3d/%3d' % (i + 1, num_samples),
np.mean(np.argmax(ens_predictions, axis=1) == targets),
nll(ens_predictions / (i + 1), targets)
]
table = tabulate.tabulate([values],
columns,
tablefmt='simple',
floatfmt='8.4f')
if i == 0:
cov_factor = np.vstack((u[None, :], v[None, :]))
subspace = SubspaceModel(torch.FloatTensor(mean),
torch.FloatTensor(cov_factor))
coords = np.dot(
cov_factor / np.sum(np.square(cov_factor), axis=1, keepdims=True),
(w - mean[None, :]).T).T
theta = torch.FloatTensor(coords[2, :])
for i in range(3):
v = subspace(torch.FloatTensor(coords[i]))
offset = 0
for param in model.parameters():
param.data.copy_(v[offset:offset + param.numel()].view(
param.size()).to(args.device))
offset += param.numel()
utils.bn_update(loaders_bn['train'], model)
print("Performance of model", i, "on the curve", end=":")
print(utils.eval(loaders['test'], model, losses.cross_entropy))
else:
assert len(args.checkpoint) == 1
swag_model = SWAG(model_cfg.base,
num_classes=num_classes,
subspace_type='pca',
subspace_kwargs={
'max_rank': 20,
'pca_rank': args.rank,
},
*model_cfg.args,
**model_cfg.kwargs)
swag_model.to(args.device)
#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 = sgd_preds.copy()
else:
#TODO: rewrite in a numerically stable way
sgd_ens_preds = sgd_ens_preds * n_ensembled / (
n_ensembled + 1) + sgd_preds / (n_ensembled + 1)
n_ensembled += 1
swag_model.collect_model(model)
if epoch == 0 or epoch % args.eval_freq == args.eval_freq - 1 or epoch == args.epochs - 1:
swag_model.set_swa()
utils.bn_update(loaders['train'], swag_model)
swag_res = utils.eval(loaders['test'], swag_model, criterion)
else:
swag_res = {'loss': None, 'accuracy': None}
if (epoch + 1) % args.save_freq == 0:
utils.save_checkpoint(args.dir,
epoch + 1,
state_dict=model.state_dict(),
optimizer=optimizer.state_dict())
if args.swag:
utils.save_checkpoint(
args.dir,
epoch + 1,
name='swag',
state_dict=swag_model.state_dict(),