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)
