| import torch
|
| from sklearn.metrics import classification_report, f1_score, accuracy_score
|
| import wandb
|
| import time
|
| import copy
|
| import torchvision
|
| import matplotlib.pyplot as plt
|
| import numpy as np
|
|
|
| class Trainer:
|
| def __init__(self, method, model, criterion, optimizer, dataloaders, target_names, output_folders, num_epochs):
|
| 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.cuda(), labels.long().cuda()
|
| 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
|
|
|
| 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.cuda(), labels_v.long().cuda()
|
| 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
|
|
|
| def main_loop(self):
|
| training_losses = []
|
| for epoch in range(self.num_epochs):
|
| start = time.time()
|
|
|
| train_loss_all = self.training_loop()
|
| train_loss_avg = (sum(train_loss_all) / len(train_loss_all))
|
| training_losses.append(train_loss_avg)
|
| wandb.log({
|
| "Train Loss": train_loss_avg})
|
| end = time.time()
|
| print('Epoch: {}. '
|
| 'Time: {:.2f} '
|
| 'train_loss: {:.4f} '.format(epoch,
|
| end - start,
|
| train_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
|
|
|
| state = {
|
| 'model': self.model.state_dict(),
|
| 'optimizer': self.optimizer.state_dict(),
|
| 'epoch': epoch,
|
| }
|
| torch.save(state, self.output_folders + 'last_epoch.pth')
|
| return
|
|
|
|
|
| class SupervisedTrainer:
|
| def __init__(self, device, method, model, classifier, criterion, optimizer, dataloaders, target_names,
|
| output_folders, num_epochs):
|
| self.device = device
|
| self.method = method
|
| self.model = model
|
| self.classifier = classifier
|
| 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 = []
|
| labels_true_train = []
|
| labels_predicted_train = []
|
| self.model.eval()
|
| self.classifier.train()
|
|
|
| for i, (inputs, labels) in enumerate(self.dataloaders['train']):
|
| inputs, labels = inputs.to(self.device), labels.long().to(self.device)
|
| self.optimizer.zero_grad()
|
| bsz = labels.shape[0]
|
|
|
| with torch.no_grad():
|
| if self.method =='SupConEF':
|
| features = self.model.calculate_embedding(inputs)
|
| elif self.method =='SupCon':
|
| features = self.model.encoder(inputs)
|
| train_output = self.classifier(features.detach())
|
| train_loss = self.criterion(train_output, labels)
|
| _, train_predicted = torch.max(train_output.data, 1)
|
| labels_true_train.append(labels.tolist())
|
| labels_predicted_train.append(train_predicted.tolist())
|
|
|
| train_loss_all.append(train_loss.cpu().data.item())
|
|
|
| train_loss.backward()
|
| self.optimizer.step()
|
|
|
| true_train = sum(labels_true_train, [])
|
| pred_train = sum(labels_predicted_train, [])
|
| train_f1 = f1_score(true_train, pred_train, average='macro')
|
| train_accuracy = accuracy_score(true_train, pred_train)
|
|
|
| return train_accuracy, train_f1, train_loss_all
|
|
|
| def validation_loop(self):
|
| validation_loss_all = []
|
| labels_true = []
|
| labels_predicted = []
|
| self.model.eval()
|
| self.classifier.eval()
|
| with torch.no_grad():
|
| for inputs_v, labels_v in self.dataloaders['valid']:
|
| inputs_v, labels_v = inputs_v.to(self.device), labels_v.long().to(self.device)
|
|
|
| bsz_v = labels_v.shape[0]
|
| if self.method =='SupConEF':
|
| outputs_v = self.classifier(self.model.calculate_embedding(inputs_v))
|
| elif self.method =='SupCon':
|
| outputs_v = self.classifier(self.model.encoder(inputs_v))
|
|
|
|
|
|
|
| validation_loss = self.criterion(outputs_v, labels_v)
|
| _, predicted = torch.max(outputs_v.data, 1)
|
| labels_true.append(labels_v.tolist())
|
| labels_predicted.append(predicted.tolist())
|
|
|
| validation_loss_all.append(validation_loss.cpu().data.item())
|
| true = sum(labels_true, [])
|
| pred = sum(labels_predicted, [])
|
| valid_f1 = f1_score(true, pred, average='macro')
|
| valid_accuracy = accuracy_score(true, pred)
|
|
|
| return valid_accuracy, valid_f1, true, pred, validation_loss_all
|
|
|
| def main_loop(self):
|
| validation_losses = []
|
| valids = []
|
| f1s = []
|
| for epoch in range(self.num_epochs):
|
| start = time.time()
|
|
|
| train_accuracy, train_f1, train_loss_all = self.training_loop()
|
| valid_accuracy, valid_f1, true, pred, validation_loss_all = self.validation_loop()
|
| validation_loss_avg = (sum(validation_loss_all) / len(validation_loss_all))
|
| train_loss_avg = (sum(train_loss_all) / len(train_loss_all))
|
| validation_losses.append(validation_loss_avg)
|
| valids.append(valid_accuracy)
|
| f1s.append(valid_f1)
|
|
|
| wandb.log({
|
| "Train Loss": train_loss_avg,
|
| "Validation Loss": validation_loss_avg,
|
| "Validation Accuracy": valid_accuracy,
|
| "Validation F1": valid_f1})
|
|
|
| print('Epoch: {}. '
|
| 'train_loss: {:.4f} '
|
| 'train_accuracy: {:.4f} '
|
| 'validation_loss: {:.4f} '
|
| 'validation_accuracy: {:.4f} '
|
| 'valid_f1:{:.4f}'.format(epoch,
|
| train_loss_avg,
|
| train_accuracy,
|
| validation_loss_avg,
|
| valid_accuracy,
|
| valid_f1))
|
|
|
| if validation_loss_avg <= min(validation_losses):
|
| torch.save(copy.deepcopy(self.model.state_dict()),
|
| self.output_folders + 'least_validation_loss_epoch.pth')
|
| torch.save(copy.deepcopy(self.classifier.state_dict()),
|
| self.output_folders + 'least_validation_loss_epoch_classifier.pth')
|
|
|
| print('Minimum Validation Loss')
|
| print(classification_report(true, pred, target_names=self.target_names))
|
| wandb.run.summary["Least Validation Loss"] = validation_loss_avg
|
|
|
| if valid_accuracy >= max(valids):
|
| torch.save(copy.deepcopy(self.model.state_dict()), self.output_folders + 'max_accuracy_epoch.pth')
|
| torch.save(copy.deepcopy(self.classifier.state_dict()),
|
| self.output_folders + 'max_accuracy_epoch_classifier.pth')
|
|
|
| print('Maximum Accuracy')
|
| print(classification_report(true, pred, target_names=self.target_names))
|
| wandb.run.summary["Best Validation Accuracy"] = valid_accuracy
|
|
|
| if valid_f1 >= max(f1s):
|
| torch.save(copy.deepcopy(self.model.state_dict()), self.output_folders + 'max_f1_epoch.pth')
|
| torch.save(copy.deepcopy(self.classifier.state_dict()),
|
| self.output_folders + 'max_f1_epoch_classifier.pth')
|
| print('Maximum F1')
|
| print(classification_report(true, pred, target_names=self.target_names))
|
| wandb.run.summary["Best Validation F1"] = valid_f1
|
| end = time.time()
|
| print('Total time for One Epoch : {:.4f} Seconds ', end - start)
|
| print('')
|
| return
|
|
|
|
|
| class FullySupervisedTrainer:
|
| 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 = []
|
| labels_true_train = []
|
| labels_predicted_train = []
|
| self.model.train()
|
|
|
| for i, (inputs, labels) in enumerate(self.dataloaders['train']):
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| inputs, labels = inputs.to(self.device), labels.long().to(self.device)
|
|
|
| self.optimizer.zero_grad()
|
| train_output = self.model(inputs)
|
| train_loss = self.criterion(train_output, labels)
|
| _, train_predicted = torch.max(train_output.data, 1)
|
|
|
|
|
|
|
| labels_true_train.append(labels.tolist())
|
| labels_predicted_train.append(train_predicted.tolist())
|
|
|
| train_loss_all.append(train_loss.cpu().data.item())
|
|
|
| train_loss.backward()
|
| self.optimizer.step()
|
|
|
| true_train = sum(labels_true_train, [])
|
| pred_train = sum(labels_predicted_train, [])
|
| train_f1 = f1_score(true_train, pred_train, average='macro')
|
| train_accuracy = accuracy_score(true_train, pred_train)
|
|
|
| return train_accuracy, train_f1, train_loss_all
|
|
|
| def validation_loop(self):
|
| validation_loss_all = []
|
| labels_true = []
|
| labels_predicted = []
|
| self.model.eval()
|
| with torch.no_grad():
|
|
|
| for inputs_v, labels_v in self.dataloaders['valid']:
|
|
|
|
|
|
|
|
|
|
|
|
|
| inputs_v, labels_v = inputs_v.to(self.device), labels_v.long().to(self.device)
|
| outputs_v = self.model(inputs_v)
|
| validation_loss = self.criterion(outputs_v, labels_v)
|
| _, predicted = torch.max(outputs_v.data, 1)
|
|
|
|
|
|
|
| labels_true.append(labels_v.tolist())
|
| labels_predicted.append(predicted.tolist())
|
| validation_loss_all.append(validation_loss.cpu().data.item())
|
| true = sum(labels_true, [])
|
| pred = sum(labels_predicted, [])
|
| valid_f1 = f1_score(true, pred, average='macro')
|
| valid_accuracy = accuracy_score(true, pred)
|
|
|
| return valid_accuracy, valid_f1, true, pred, validation_loss_all
|
|
|
| def main_loop(self):
|
| validation_losses = []
|
| valids = []
|
| f1s = []
|
| for epoch in range(self.num_epochs):
|
| start = time.time()
|
|
|
| train_accuracy, train_f1, train_loss_all = self.training_loop()
|
| valid_accuracy, valid_f1, true, pred, validation_loss_all = self.validation_loop()
|
| validation_loss_avg = (sum(validation_loss_all) / len(validation_loss_all))
|
| train_loss_avg = (sum(train_loss_all) / len(train_loss_all))
|
| validation_losses.append(validation_loss_avg)
|
| valids.append(valid_accuracy)
|
| f1s.append(valid_f1)
|
|
|
| wandb.log({
|
| "Train Loss": train_loss_avg,
|
| "Train Accuracy": train_accuracy,
|
| "Validation Loss": validation_loss_avg,
|
| "Validation Accuracy": valid_accuracy,
|
| "Validation F1": valid_f1})
|
|
|
| print('Epoch: {}. '
|
| 'train_loss: {:.4f} '
|
| 'train_accuracy: {:.4f} '
|
| 'validation_loss: {:.4f} '
|
| 'validation_accuracy: {:.4f} '
|
| 'valid_f1:{:.4f}'.format(epoch,
|
| train_loss_avg,
|
| train_accuracy,
|
| validation_loss_avg,
|
| valid_accuracy,
|
| valid_f1))
|
|
|
| if validation_loss_avg <= min(validation_losses):
|
| torch.save(copy.deepcopy(self.model.state_dict()),
|
| self.output_folders + 'least_validation_loss_epoch.pth')
|
|
|
| print('Minimum Validation Loss')
|
| print(classification_report(true, pred, target_names=self.target_names))
|
| wandb.run.summary["Least Validation Loss"] = validation_loss_avg
|
|
|
| if valid_accuracy >= max(valids):
|
| torch.save(copy.deepcopy(self.model.state_dict()), self.output_folders + 'max_accuracy_epoch.pth')
|
|
|
| print('Maximum Accuracy')
|
| print(classification_report(true, pred, target_names=self.target_names))
|
| wandb.run.summary["Best Validation Accuracy"] = valid_accuracy
|
|
|
| if valid_f1 >= max(f1s):
|
| torch.save(copy.deepcopy(self.model.state_dict()), self.output_folders + 'max_f1_epoch.pth')
|
| print('Maximum F1')
|
| print(classification_report(true, pred, target_names=self.target_names))
|
| wandb.run.summary["Best Validation F1"] = valid_f1
|
| end = time.time()
|
| print('Total time for One Epoch : {:.4f} Seconds ', end - start)
|
| print('')
|
| return
|
|
|
