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