"""
Created on December 11, 2021.
Train_Valid_ACL.py
@author: Soroosh Tayebi Arasteh <soroosh.arasteh@fau.de>
https://github.com/tayebiarasteh/
"""
import os.path
import time
import pdb
from tensorboardX import SummaryWriter
import torch
import torchmetrics
import torchio as tio
from config.serde import read_config, write_config
import warnings
warnings.filterwarnings('ignore')
class Training:
def __init__(self, cfg_path, num_iterations=100, resume=False, torch_seed=None):
"""This class represents training and validation processes.
Parameters
----------
cfg_path: str
Config file path of the experiment
num_iterations: int
Total number of epochs for training
resume: bool
if we are resuming training from a checkpoint
torch_seed: int
Seed used for random generators in PyTorch functions
"""
self.params = read_config(cfg_path)
self.cfg_path = cfg_path
self.num_iterations = num_iterations
if resume == False:
self.model_info = self.params['Network']
self.model_info['seed'] = torch_seed or self.model_info['seed']
self.iteration = 0
self.best_F1 = float('inf')
self.setup_cuda()
self.writer = SummaryWriter(log_dir=os.path.join(self.params['target_dir'], self.params['tb_logs_path']))
def setup_cuda(self, cuda_device_id=0):
"""setup the device.
Parameters
----------
cuda_device_id: int
cuda device id
"""
if torch.cuda.is_available():
torch.backends.cudnn.fastest = True
torch.cuda.set_device(cuda_device_id)
self.device = torch.device('cuda')
torch.cuda.manual_seed_all(self.model_info['seed'])
torch.manual_seed(self.model_info['seed'])
else:
self.device = torch.device('cpu')
def time_duration(self, start_time, end_time):
"""calculating the duration of training or one iteration
Parameters
----------
start_time: float
starting time of the operation
end_time: float
ending time of the operation
Returns
-------
elapsed_hours: int
total hours part of the elapsed time
elapsed_mins: int
total minutes part of the elapsed time
elapsed_secs: int
total seconds part of the elapsed time
"""
elapsed_time = end_time - start_time
elapsed_hours = int(elapsed_time / 3600)
if elapsed_hours >= 1:
elapsed_mins = int((elapsed_time / 60) - (elapsed_hours * 60))
elapsed_secs = int(elapsed_time - (elapsed_hours * 3600) - (elapsed_mins * 60))
else:
elapsed_mins = int(elapsed_time / 60)
elapsed_secs = int(elapsed_time - (elapsed_mins * 60))
return elapsed_hours, elapsed_mins, elapsed_secs
def setup_model(self, model, optimiser, loss_function, weight=None):
"""Setting up all the models, optimizers, and loss functions.
Parameters
----------
model: model file
The network
optimiser: optimizer file
The optimizer
loss_function: loss file
The loss function
weight: 1D tensor of float
class weights
"""
# prints the network's total number of trainable parameters and
# stores it to the experiment config
total_param_num = sum(p.numel() for p in model.parameters() if p.requires_grad)
print(f'\nTotal # of trainable parameters: {total_param_num:,}')
print('----------------------------------------------------\n')
self.model = model.to(self.device)
if not weight==None:
self.loss_weight = weight.to(self.device)
self.loss_function = loss_function(weight=self.loss_weight)
else:
self.loss_function = loss_function()
self.optimiser = optimiser
self.model_info['total_param_num'] = total_param_num
self.model_info['loss_function'] = loss_function.__name__
self.model_info['num_iterations'] = self.num_iterations
self.params['Network'] = self.model_info
write_config(self.params, self.cfg_path, sort_keys=True)
def load_checkpoint(self, model, optimiser, loss_function):
"""In case of resuming training from a checkpoint,
loads the weights for all the models, optimizers, and
loss functions, and device, tensorboard events, number
of iterations (epochs), and every info from checkpoint.
Parameters
----------
model: model file
The network
optimiser: optimizer file
The optimizer
loss_function: loss file
The loss function
"""
checkpoint = torch.load(os.path.join(self.params['target_dir'], self.params['network_output_path']) + '/' + self.params['checkpoint_name'])
self.device = None
self.model_info = checkpoint['model_info']
self.setup_cuda()
self.model = model.to(self.device)
self.loss_weight = checkpoint['loss_state_dict']['weight']
self.loss_weight = self.loss_weight.to(self.device)
self.loss_function = loss_function(weight=self.loss_weight)
self.optimiser = optimiser
self.model.load_state_dict(checkpoint['model_state_dict'])
self.optimiser.load_state_dict(checkpoint['optimizer_state_dict'])
self.iteration = checkpoint['iteration']
self.best_F1 = checkpoint['best_F1']
self.writer = SummaryWriter(log_dir=os.path.join(os.path.join(
self.params['target_dir'], self.params['tb_logs_path'])), purge_step=self.iteration + 1)
def execute_training(self, train_loader, valid_loader=None, augmentation=False):
"""Executes training by running training and validation at each epoch.
Parameters
----------
train_loader: Pytorch dataloader object
training data loader
valid_loader: Pytorch dataloader object
validation data loader
"""
self.params = read_config(self.cfg_path)
total_start_time = time.time()
for iteration in range(self.num_iterations - self.iteration):
self.iteration += 1
start_time = time.time()
train_F1, train_acc, train_loss = self.train_epoch_3D(train_loader=train_loader)
if not valid_loader == None:
valid_F1, valid_acc, valid_loss = self.valid_epoch_3D(valid_loader=valid_loader)
# Validation iteration & calculate metrics
if (self.iteration) % (self.params['epochbased_save_freq']) == 0:
end_time = time.time()
iteration_hours, iteration_mins, iteration_secs = self.time_duration(start_time, end_time)
total_hours, total_mins, total_secs = self.time_duration(total_start_time, end_time)
# saving the model, checkpoint, TensorBoard, etc.
if not valid_loader == None:
self.calculate_tb_stats(train_F1=train_F1, train_acc=train_acc, train_loss=train_loss,
valid_F1=valid_F1, valid_acc=valid_acc, valid_loss=valid_loss)
self.savings_prints(iteration_hours, iteration_mins, iteration_secs, total_hours,
total_mins, total_secs, train_F1, train_acc, train_loss,
valid_F1, valid_acc, valid_loss)
else:
self.calculate_tb_stats(train_F1=train_F1, train_acc=train_acc, train_loss=train_loss)
self.savings_prints(iteration_hours, iteration_mins, iteration_secs, total_hours,
total_mins, total_secs, train_F1, train_acc, train_loss)
def train_epoch_3D(self, train_loader):
"""This is the pipeline based on Pytorch's Dataset and Dataloader
Parameters
----------
train_loader: Pytorch dataloader object
training data loader
Returns
-------
average_f1_score: float
average training F1 score of the epoch
average_accuracy: float
average training accuracy of the epoch
average_loss: float
average training loss of the epoch
"""
self.model.train()
total_loss = 0.0
total_accuracy = 0.0
total_f1_score = 0.0
# we imagine we only have one batch
image = train_loader
label = torch.ones((1, 1))
label = label.long()
image = image.float()
image = image.to(self.device)
label = label.to(self.device)
self.optimiser.zero_grad()
with torch.autograd.set_detect_anomaly(True):
output, a_output = self.model(image)
max_a_output = a_output.argmax(dim=2) # get the slice with ACL
loss = self.loss_function(output, label[:, 0])
loss.backward()
self.optimiser.step()
total_loss += loss.item()
# TODO: evaluation metric calculation
return average_f1_score, average_accuracy, average_loss
def savings_prints(self, iteration_hours, iteration_mins, iteration_secs,
total_hours, total_mins, total_secs, train_F1, train_acc,
train_loss, valid_F1=None, valid_acc=None, valid_loss=None):
"""Saving the model weights, checkpoint, information,
and training and validation loss and evaluation statistics.
Parameters
----------
iteration_hours: int
hours part of the elapsed time of each iteration
iteration_mins: int
minutes part of the elapsed time of each iteration
iteration_secs: int
seconds part of the elapsed time of each iteration
total_hours: int
hours part of the total elapsed time
total_mins: int
minutes part of the total elapsed time
total_secs: int
seconds part of the total elapsed time
train_loss: float
training loss of the model
valid_loss: float
validation loss of the model
train_acc: float
training accuracy of the model
valid_acc: float
validation accuracy of the model
train_F1: float
training F1 score of the model
valid_F1: float
validation F1 score of the model
"""
# Saves information about training to config file
self.params['Network']['num_steps'] = self.iteration
write_config(self.params, self.cfg_path, sort_keys=True)
# Saving the model based on the best F1
if valid_F1:
if valid_F1 < self.best_F1:
self.best_F1 = valid_F1
torch.save(self.model.state_dict(), os.path.join(self.params['target_dir'], self.params['network_output_path']) + '/' +
self.params['trained_model_name'])
else:
if train_F1 < self.best_F1:
self.best_F1 = train_F1
torch.save(self.model.state_dict(), os.path.join(self.params['target_dir'], self.params['network_output_path']) + '/' +
self.params['trained_model_name'])
# Saving every couple of iterations
if (self.iteration) % self.params['network_save_freq'] == 0:
torch.save(self.model.state_dict(), os.path.join(self.params['target_dir'], self.params['network_output_path']) + '/' +
'iteration{}_'.format(self.iteration) + self.params['trained_model_name'])
# Save a checkpoint every 2 iterations
if (self.iteration) % self.params['network_checkpoint_freq'] == 0:
torch.save({'iteration': self.iteration,
'model_state_dict': self.model.state_dict(),
'optimizer_state_dict': self.optimiser.state_dict(),
'loss_state_dict': self.loss_function.state_dict(), 'num_iterations': self.num_iterations,
'model_info': self.model_info, 'best_F1': self.best_F1},
os.path.join(self.params['target_dir'], self.params['network_output_path']) + '/' + self.params['checkpoint_name'])
print('------------------------------------------------------'
'----------------------------------')
print(f'Iteration: {self.iteration}/{self.num_iterations} | '
f'Iteration Time: {iteration_hours}h {iteration_mins}m {iteration_secs}s | '
f'Total Time: {total_hours}h {total_mins}m {total_secs}s')
print(f'\n\tTrain Loss: {train_loss:.4f} | Acc: {train_acc * 100:.2f}% | F1: {train_F1 * 100:.2f}%')
if valid_loss:
print(f'\t Val. Loss: {valid_loss:.4f} | Acc: {valid_acc * 100:.2f}% | F1: {valid_F1 * 100:.2f}%')
# saving the training and validation stats
msg = f'----------------------------------------------------------------------------------------\n' \
f'Iteration: {self.iteration}/{self.num_iterations} | Iteration Time: {iteration_hours}h {iteration_mins}m {iteration_secs}s' \
f' | Total Time: {total_hours}h {total_mins}m {total_secs}s\n\n\tTrain Loss: {train_loss:.4f} | ' \
f'Acc: {train_acc * 100:.2f}% | ' \
f'F1: {train_F1 * 100:.2f}%\n\t Val. Loss: {valid_loss:.4f} | Acc: {valid_acc*100:.2f}% | F1: {valid_F1 * 100:.2f}%\n\n'
else:
msg = f'----------------------------------------------------------------------------------------\n' \
f'Iteration: {self.iteration}/{self.num_iterations} | Iteration Time: {iteration_hours}h {iteration_mins}m {iteration_secs}s' \
f' | Total Time: {total_hours}h {total_mins}m {total_secs}s\n\n\tTrain Loss: {train_loss:.4f} | ' \
f'Acc: {train_acc * 100:.2f}% | F1: {train_F1 * 100:.2f}%\n\n'
with open(os.path.join(self.params['target_dir'], self.params['stat_log_path']) + '/Stats', 'a') as f:
f.write(msg)
def calculate_tb_stats(self, train_F1, train_acc, train_loss, valid_F1=None, valid_acc=None, valid_loss=None):
"""Adds the evaluation metrics and loss values to the tensorboard.
Parameters
----------
train_loss: float
training loss of the model
valid_loss: float
validation loss of the model
train_acc: float
training accuracy of the model
valid_acc: float
validation accuracy of the model
train_F1: float
training F1 score of the model
valid_F1: float
validation F1 score of the model
"""
self.writer.add_scalar('Train_F1', train_F1, self.iteration)
# self.writer.add_scalar('Train_Accuracy', train_acc, self.iteration)
self.writer.add_scalar('Train_Loss', train_loss, self.iteration)
if valid_F1 is not None:
self.writer.add_scalar('Valid_F1', valid_F1, self.iteration)
# self.writer.add_scalar('Valid_Accuracy', valid_acc, self.iteration)
self.writer.add_scalar('Valid_Loss', valid_loss, self.iteration)
Datasets
Models
