optimizer = torch.optim.SGD(model.parameters(), 0.1, momentum=0.9, weight_decay=1e-4) # optimiser = optim.RMSprop(model.parameters(), alpha=0.9, lr=0.0001, weight_decay=1e-6) loss_function = nn.CrossEntropyLoss() # device = "cuda:0" if torch.cuda.is_available() else "cpu" trial = Trial(model, optimizer, loss_function, metrics=['loss', 'acc', 'top_5_acc'], callbacks=[ callbacks.TensorBoard(write_graph=False, comment=f'resnet50_{n_bn}_{rep}'), callbacks.MultiStepLR([30, 60]), callbacks.MostRecent(dir + model_file + '_{epoch:02d}.pt') ]).to('cuda') trial.with_generators(trainloader, test_generator=testloader) pattern = re.compile(model_file + '_\d+.pt') for filepath in os.listdir(dir): if pattern.match(filepath): trial.load_state_dict(torch.load(dir + filepath)) trial.run(epochs=90) trial.evaluate(data_key=torchbearer.TEST_DATA) torch.save(model.module.state_dict(), dir + model_file + '.pt')
if mode == 'fmix': aug = [FMix(alpha=1, decay_power=3)] model = VAE(64, var=args.var) trial = Trial(model, optim.Adam(model.parameters(), lr=5e-2), nll, metrics=[ metrics.MeanSquaredError(pred_key=SAMPLE), metrics.mean(NLL), metrics.mean(KL), 'loss' ], callbacks=[ sample, kld(distributions.Normal(0, 1)), init.XavierNormal(targets=['Conv']), callbacks.MostRecent(args.dir + '/' + mode + '_' + str(args.i) + '.pt'), callbacks.MultiStepLR([40, 80]), callbacks.TensorBoard(write_graph=False, comment=mode + '_' + str(args.i), log_dir='vae_logs'), *aug ]) if mode in ['base', 'mix', 'cutmix']: trial = trial.load_state_dict( torch.load('vaes/' + '/' + mode + '_' + str(args.i) + '.pt')) trial.with_generators(train_loader, test_loader).to('cuda').run(args.epochs, verbose=1)
super().__init__('roc_auc_score', lambda y_pred, y_true: sklearn.metrics.roc_auc_score(y_true.cpu().numpy(), y_pred.detach().sigmoid().cpu().numpy()), running=False) metrics_append = [RocAucScore()] if args.dataset == 'imagenet_a': from datasets.imagenet_a import indices_in_1k @on_forward_validation def map_preds(state): state[torchbearer.PREDICTION] = state[torchbearer.PREDICTION][:, indices_in_1k] cb.append(map_preds) print('==> Training model..') acc = 'acc' if args.dataset in ['toxic', 'toxic_bert', 'imdb', 'yelp_2']: acc = 'binary_acc' trial = Trial(net, optimizer, criterion, metrics=[acc, 'loss', 'lr'] + metrics_append, callbacks=cb) trial.with_generators(train_generator=trainloader, val_generator=valloader, train_steps=args.train_steps, test_generator=testloader).to(args.device) if args.reload: state = torch.load(args.model_file) trial.load_state_dict(state, resume=args.dataset != 'imagenet_a') trial.replay() if trainloader is not None: trial.run(args.epoch, verbose=args.verbose) trial.evaluate(data_key=torchbearer.TEST_DATA)
CosineAnnealingLR(args.epoch, eta_min=0.)) cb.append(WarmupLR(0.1, args.lr)) if args.lr_warmup else [] # FMix loss is equivalent to mixup loss and works for all msda in torchbearer if args.msda_mode not in [None, 'None']: bce = True if args.dataset == 'toxic' else False criterion = modes['fmix'].loss(bce) elif args.dataset == 'toxic': criterion = nn.BCEWithLogitsLoss() else: criterion = nn.CrossEntropyLoss() # from torchbearer.metrics.roc_auc_score import RocAucScore print('==> Training model..') trial = Trial(net, optimizer, criterion, metrics=['acc', 'loss', 'lr'], callbacks=cb) trial.with_generators(train_generator=trainloader, val_generator=valloader, train_steps=args.train_steps, test_generator=testloader).to(args.device) if args.reload: state = torch.load(args.model_file) trial.load_state_dict(state) trial.replay() trial.run(args.epoch, verbose=args.verbose) trial.evaluate(data_key=torchbearer.TEST_DATA)
def train(): # train_loader = DataLoader(train_dataset, batch_size=batch_size, shuffle=True) print('start training') dataset_size = len(dataset) indices = list(range(dataset_size)) split = int(np.floor(validation_split * dataset_size)) if shuffle_dataset: np.random.seed(random_seed) np.random.shuffle(indices) train_indices, val_indices = indices[split:], indices[:split] # Creating PT data samplers and loaders: train_sampler = SubsetRandomSampler(train_indices) valid_sampler = SubsetRandomSampler(val_indices) train_loader = torch.utils.data.DataLoader(dataset, batch_size=batch_size, sampler=train_sampler) val_loader = torch.utils.data.DataLoader(dataset, batch_size=batch_size, sampler=valid_sampler) model = resnet50(pretrained=True) model.avgpool = nn.AdaptiveAvgPool2d((1, 1)) model.fc = nn.Linear(2048, len(dataset.classes)) model.train() # Freeze layers by not tracking gradients for param in model.parameters(): param.requires_grad = False model.fc.weight.requires_grad = True #unfreeze last layer weights model.fc.bias.requires_grad = True #unfreeze last layer biases optimiser1 = optim.Adam(filter(lambda p: p.requires_grad, model.parameters()), lr=1e-3) #only optimse non-frozen layers loss_function = nn.CrossEntropyLoss() device = "cuda:0" if torch.cuda.is_available() else "cpu" print(device) checkpointer = torchbearer.callbacks.checkpointers.Best( filepath='model.pt', monitor='loss') trial = Trial(model, optimiser1, loss_function, metrics=['loss', 'accuracy'], callbacks=[checkpointer]).to(device) trial.with_generators(train_loader, val_generator=val_loader) trial.run(epochs=10) state_dict = torch.load('model.pt') trial_reloaded = Trial(model, optimiser1, loss_function, metrics=['loss', 'accuracy'], callbacks=[checkpointer]).to(device) trial_reloaded.with_generators(train_loader, val_generator=val_loader) trial_reloaded.load_state_dict(state_dict) trial_reloaded.run(epochs=20) results = trial_reloaded.evaluate(data_key=torchbearer.VALIDATION_DATA) print() print(results) torch.save(model, 'mymodel.pth')
class eval_model(nn.Module): def __init__(self, model, loss_f, device, metrics, trial_name, opt=optim.Adam,opt_params={}, **kwargs): super().__init__() self.model = model self.loss_f = loss_f self.opt = opt(list(self.model.parameters()), **opt_params) self.metrics = metrics self.device = device self.trial_path = os.path.join('logs/', trial_name) self.call_backs = None self.best_step = 0 def init_trial(self): #initial call_back functions for the trail. if self.call_backs: self.trial = Trial(self.model,self.opt,self.loss_f,metrics=self.metrics,callbacks=self.call_backs).to(self.device) else: self.trial = Trial(self.model,self.opt,self.loss_f,metrics=self.metrics).to(self.device) def run(self, epoch, tr_loader,val_loader,t_loader=None,val_steps=None): self.trial.with_generators(tr_loader, test_generator=t_loader,val_generator=val_loader,val_steps=val_steps) return self.trial.run(epoch) def save_model(self, suffix, path=None, mkdir=True, **kwargs): if path is None: path = os.path.join(self.trial_path, "model_{}.mf".format(suffix)) if mkdir: os.makedirs(os.path.dirname(path), exist_ok=True) torch.save(OrderedDict([ ('model', self.model.state_dict(**kwargs)), ('trial', self.trial.state_dict()), ('opt', self.opt.state_dict()), ('best_step', self.best_step), ]), path) return path def load_model(self, suffix, path=None, **kwargs): if path is None: path = os.path.join(self.trial_path, "model_{}.mf".format(suffix)) checkpoint = torch.load(path) self.trial.load_state_dict(checkpoint['trial']) self.model.load_state_dict(checkpoint['model'], **kwargs) self.opt.load_state_dict(checkpoint['opt']) self.best_step = int(checkpoint['best_step']) def plot_loss(self, loss, title, xlabel, x_all=False): plt.figure() ax = plt.axes() ax.set_title(title) ax.set_ylabel("loss") ax.set_xlabel(xlabel) x = range(len(loss)) plt.plot(x,loss) if x_all: ax.xaxis.set_major_locator(MaxNLocator(integer=True)) plt.xticks(x) plt.legend() plt.show() def plot_fig(self,out,epoch,title,y_label, para=['loss','val_loss']): tem = [ [0]*epoch for i in range(len(para))] for idx, label in enumerate(para): for i in range(epoch): tem[idx][i] = out[i][label] plt.figure() ax = plt.axes() ax.set_title(title) ax.set_ylabel(y_label) ax.set_xlabel("Epoch") for i, item in enumerate(tem): plt.plot(range(epoch),item, label=para[i]) plt.legend() plt.show()
optimizer=optimizer, criterion=F.binary_cross_entropy, metrics=['loss'], callbacks=[checkpointer]) torchbearer_trial.with_train_generator(traingen) torchbearer_trial.run(epochs=3) # Reloading the trial state_dict = torch.load('model.pt') model = BasicModel() trial_reloaded = Trial(model, optimizer=optimizer, criterion=F.binary_cross_entropy, metrics=['loss'], callbacks=[checkpointer]) trial_reloaded.load_state_dict(state_dict) trial_reloaded.with_train_generator(traingen) trial_reloaded.run(epochs=6) # We cannot simply load_state_dict into a native PyTorch module since there are additional Trial attributes in there model = BasicModel() try: model.load_state_dict(state_dict) except AttributeError as e: print("\n") print(e) # Prepare the torchbearer trial once again, but this time with save_model_params_only option at checkpointer model = BasicModel() checkpointer = torchbearer.callbacks.checkpointers.Best( filepath='model.pt', monitor='loss', save_model_params_only=True)