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,