scheduler = lr_scheduler.MultiStepLR(optimizer,
args.milestones,
gamma=args.step_gamma)
ang_sb = []
ang_np = []
p_angles = []
# Modified code from here to calcuate top subspace of Hessian at every iterate
for i in range(args.num_models):
model_weights_path = osp.join(ckpt_dir,
f'model_weights_epochs_{i}.pth')
print(f'Processing model number {i}')
state_dict = torch.load(model_weights_path, device)
model.load_state_dict(state_dict, strict=True)
model = model.to(device)
C = 10
Hess = FullHessian(
crit='CrossEntropyLoss',
loader=loader,
device=device,
model=model,
num_classes=C,
hessian_type='Hessian',
init_poly_deg=
64, # number of iterations used to compute maximal/minimal eigenvalue
poly_deg=
128, # the higher the parameter the better the approximation
spectrum_margin=0.05,
poly_points=1024, # number of points in spectrum approximation
optimizer = optim.Adam(model.parameters(),
lr=args.learning_rate,
weight_decay=args.weight_decay)
else:
raise Exception('Optimizer not allowed: {}'.format(args.optimizer))
scheduler = lr_scheduler.MultiStepLR(optimizer,
args.milestones,
gamma=args.step_gamma)
if args.resume is not None:
assert osp.exists(args.resume)
assert osp.isfile(args.resume)
ckpt = torch.load(args.resume)
assert 'model' in ckpt.keys()
model.load_state_dict(ckpt['model'], strict=True)
system = train(model,
optimizer,
scheduler,
dataloaders,
criterion,
device,
args.sp,
args.sf,
num_epochs=args.num_epochs,
args=args,
ckpt_dir=ckpt_dir,
dataset_sizes=dataset_sizes,
images_dir=images_dir)
optimizer = optim.SGD(model.parameters(),
lr=args.learning_rate,
weight_decay=args.weight_decay,
momentum=0.9)
# optimizer = optim.SGD(model.parameters(), lr=args.learning_rate, weight_decay=args.weight_decay, momentum=0.9)
scheduler = lr_scheduler.MultiStepLR(optimizer,
args.milestones,
gamma=args.step_gamma)
if args.resume is not None:
assert osp.exists(args.resume)
assert osp.isfile(args.resume)
ckpt = torch.load(args.resume)
assert 'model' in ckpt.keys()
model.load_state_dict(ckpt['model'])
if args.dataset in ['FashionMNIST', 'MNIST']:
dataset_sizes = {'train': 6e4, 'test': 1e4}
else:
dataset_sizes = {'train': 5e4, 'test': 1e4}
system = train(model,
optimizer,
scheduler,
dataloaders,
criterion,
device,
num_epochs=args.num_epochs,
args=args,
dataset_sizes=dataset_sizes,