File size: 5,361 Bytes
a639402 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 | import torch
import wandb
import time
from sklearn.manifold import TSNE
import matplotlib.pyplot as plt
class Trainer:
def __init__(self, device, method, model, criterion, optimizer, dataloaders, target_names, output_folders, num_epochs):
self.device = device
self.method = method
self.model = model
self.criterion = criterion
self.optimizer = optimizer
self.dataloaders = dataloaders
self.target_names = target_names
self.output_folders = output_folders
self.num_epochs = num_epochs
def training_loop(self):
train_loss_all = []
self.optimizer.zero_grad()
for i, (inputs, labels) in enumerate(self.dataloaders['train']):
self.optimizer.zero_grad()
inputs = torch.cat([inputs[0], inputs[1]], dim=0)
inputs, labels = inputs.to(self.device), labels.long().to(self.device)
bsz = labels.shape[0]
features = self.model(inputs)
f1, f2 = torch.split(features, [bsz, bsz], dim=0)
features = torch.cat([f1.unsqueeze(1), f2.unsqueeze(1)], dim=1)
if self.method == 'SupCon':
train_loss = self.criterion(features, labels)
elif self.method == 'SimCLR':
train_loss = self.criterion(features)
else:
raise ValueError('contrastive method not supported: {}'.
format(self.method))
train_loss_all.append(train_loss.cpu().data.item())
train_loss.backward()
self.optimizer.step()
return train_loss_all, features
def validation_loop(self):
validation_loss_all = []
self.model.eval()
with torch.no_grad():
for inputs_v, labels_v in self.dataloaders['valid']:
inputs_v = torch.cat([inputs_v[0], inputs_v[1]], dim=0)
inputs_v, labels_v = inputs_v.to(self.device), labels_v.long().to(self.device)
bsz_v = labels_v.shape[0]
features_v = self.model(inputs_v)
f1_v, f2_v = torch.split(features_v, [bsz_v, bsz_v], dim=0)
features_v = torch.cat([f1_v.unsqueeze(1), f2_v.unsqueeze(1)], dim=1)
if self.method == 'SupCon':
validation_loss = self.criterion(features_v, labels_v)
elif self.method == 'SimCLR':
validation_loss = self.criterion(features_v)
else:
raise ValueError('contrastive method not supported: {}'.
format(self.method))
validation_loss_all.append(validation_loss.cpu().data.item())
return validation_loss_all, features_v
def main_loop(self):
training_losses = []
valid_losses = []
for epoch in range(self.num_epochs):
start = time.time()
train_loss_all, features_train = self.training_loop()
valid_loss_all, features_v = self.validation_loop()
#tsne = TSNE(n_components=3)
#tsne_features = tsne.fit_transform(features_v.detach().cpu().numpy())
#plt.scatter(tsne_features[:,0], tsne_features[:,1])
#plt.savefig(self.output_folders + 'tsne'+str(epoch)+'.png')
train_loss_avg = (sum(train_loss_all) / len(train_loss_all))
valid_loss_avg = (sum(valid_loss_all) / len(valid_loss_all))
training_losses.append(train_loss_avg)
valid_losses.append(valid_loss_avg)
wandb.log({
"Train Loss": train_loss_avg,
"Validation Loss": valid_loss_avg})
end = time.time()
print('Epoch: {}. '
'Time: {:.2f} '
'train_loss: {:.4f} '
'valid_loss: {:.4f} '.format(epoch,
end - start,
train_loss_avg,
valid_loss_avg
))
if train_loss_avg <= min(training_losses):
state = {
'model': self.model.state_dict(),
'optimizer': self.optimizer.state_dict(),
'epoch': epoch,
}
torch.save(state, self.output_folders + 'least_training_loss_epoch.pth')
del state
wandb.run.summary["Least Train Loss"] = train_loss_avg
if valid_loss_avg <= min(valid_losses):
state = {
'model': self.model.state_dict(),
'optimizer': self.optimizer.state_dict(),
'epoch': epoch,
}
torch.save(state, self.output_folders + 'least_validation_loss_epoch.pth')
del state
wandb.run.summary["Least Valid Loss"] = valid_loss_avg
state = {
'model': self.model.state_dict(),
'optimizer': self.optimizer.state_dict(),
'epoch': epoch,
}
torch.save(state, self.output_folders + 'last_epoch.pth')
return |