scripts/train_nnmodels.py

"""
Train a neural network model with distributed training
"""


import numpy as np
import pandas as pd
import time
import argparse

import torch
import torch.nn as nn
import torch.optim as optim
from torch.utils.data import DataLoader
import torch.distributed as dist
from torch.utils.data.distributed import DistributedSampler
from torch.nn.parallel import DistributedDataParallel
import torch.multiprocessing as mp


import os
import subprocess
import importlib
import sys
from pathlib import Path

PROJECT_ROOT = Path(__file__).resolve().parents[1]
sys.path.insert(0, str(PROJECT_ROOT))
from model.ephod.training import trainutils, nn_models
importlib.reload(trainutils);




def create_parser():
    '''Parse command-line training arguments'''
    
    parser = argparse.ArgumentParser(description="Train NN mode")
    parser.add_argument('--trainseqs', type=str, 
                        help='File containing accession codes of training sequences in csv format')
    parser.add_argument('--valseqs', type=str,
                        help='File containing accession codes of validation sequences in csv format')
    parser.add_argument('--target_data', type=str,
                        help='Path to csv file containing target labels and sample weights')
    parser.add_argument('--embedding_dir', type=str,
                        help='Path to directory containing per-residue embeddings of sequences')
    parser.add_argument('--model_name', type=str, 
                        help='Name of model used in saving parameters.')                             
    parser.add_argument('--paramsjson', type=str,
                        help='Json file containing hyperparameters for building and training NN model')
    parser.add_argument('--savedir', type=str, 
                        help='Directory to save model parameters and training progress')
    parser.add_argument('--distributed', type=int,
                        help='Whether to train with a single GPU (0) or multiple GPUs(1)')
    parser.add_argument('--maxlen', default=1024, type=int,
                        help='Maximum length of sequences. Embeddings will be post-padded to this length with zeros')
    parser.add_argument('--batch_size', type=int,
                        help='Batch size of single GPU used in training and inference (i.e. local, non-distributed)')
    parser.add_argument('--num_nodes', default=1, type=int,
                        help='Number of nodes for distributed training')
    parser.add_argument('--num_workers', default=0, type=int,
                        help='Number of workers for data loading')
    parser.add_argument('--epochs', default=1000, type=int,
                        help='The maximum number of epochs for training')
    parser.add_argument('--reduce_lr_patience', default=20, type=int, 
                        help='Reduce learning rate by 0.5 if validation loss does not improve after reduce_lr_patience epochs')
    parser.add_argument('--stop_patience', default=100, type=int,
                        help='Exit training if validation loss does not improve after top_patience epochs')
    parser.add_argument('--restart', default=0, type=int,
                        help='Whether to begin new training (0) and overwrite checkpoint files, or to restart training (1) and append checkpoint files')
    parser.add_argument('--verbose', default=1, type=int,
                        help='Whether to print out details of training (1) or not (0)')
    args = parser.parse_args()
    args.params = utils.read_json(args.paramsjson)  # Model parameters as dict

    return args




def configure_cuda(args):
    '''Configure devices for training'''
    
    if not torch.cuda.is_available():
        print("FATAL: GPU not available")
        exit()
    torch.backends.cudnn.benchmark = True  # For efficient training
    args.gpus_per_node = torch.cuda.device_count()   
    args.num_cpus = os.cpu_count()     
    if args.distributed:
        args.world_size = args.gpus_per_node * args.num_nodes
    else:
        args.world_size = 1
    
    return args.world_size, args




def configure_processes(rank, args):
    '''Setup the processes group for distributed training'''
    
    os.environ['MASTER_ADDR'] = 'localhost'
    os.environ['MASTER_PORT'] = '12355'
    dist.init_process_group("nccl", rank=rank, world_size=args.world_size)




def get_dataset(seqs, sample_method, args):
    '''Get dataset of embeddings data'''

    accessions = pd.read_csv(seqs, index_col=0).iloc[:,-1].values # Seq. accession codes
    target_data = pd.read_csv(args.target_data, index_col=0)
    dataset = dataproc.EmbeddingData(
        accessions=accessions, 
        y=target_data.loc[accessions,'y'].values,
        weights=target_data.loc[accessions, sample_method].values, 
        embedding_dir=args.embedding_dir, 
        use_mask=True, 
        maxlen=args.maxlen
        )
    
    return dataset




def get_distributed_dataloader(rank, world_size, dataset, train_mode, args):
    '''Return a dataloader for distributed training/inference'''
    
    sampler = DistributedSampler(dataset, 
                                 num_replicas=world_size, 
                                 rank=rank,
                                 shuffle=train_mode, 
                                 drop_last=train_mode)
    dataloader = DataLoader(dataset, 
                            sampler=sampler,
                            batch_size=args.batch_size,
                            pin_memory=True,
                            num_workers=args.num_workers,
                            drop_last=train_mode)
    
    return dataloader    
    
    
    

def build_model(args):
    '''Build neural network model'''

    # Build model with nn.Module
    args.model_params = {
        key: args.params[key] for key in 
        ['dim', 'dim', 'kernel_size', 'dropout', 'activation', 'res_blocks', 
         'random_seed']
        }
    torch.manual_seed(args.model_params['random_seed'])  # Reproducibility
    model = nn_models.ResidualLightAttention(**args.model_params) # Replace as needed.
    args.num_parameters = torchmodels.count_parameters(model)['FULL_MODEL']
    
    return model, args




def build_optimizer(model, args):
    '''Build optimizer for model'''
    
    optimizer = optim.Adam(model.parameters(), 
                           lr=args.params['learning_rate'],
                           weight_decay=args.params['l2reg'])
    
    return optimizer
    

    

def save_model(model, optimizer, path, args):
    '''Save model and optimizer parameters to disk'''
    
    torch.save(
        {
            'model_state_dict': model.state_dict(),
            'model_params': args.model_params,
            'optimizer_state_dict': optimizer.state_dict(),
            },
        path)    




def load_model(model, path, args):
    '''Load model parameters from path'''
    
    checkpoint = torch.load(path, map_location='cpu')
    assert args.model_params == checkpoint['model_params'], \
        "Saved model params and current params are different"
    '''
    model_state_dict = {k.partition('model.')[]: v for k,v in \
                        checkpoint['model_state_dict']}
    '''
    model_state_dict = checkpoint['model_state_dict']
    model.load_state_dict(model_state_dict, strict=False)
    
    return model    
    
    


def load_optimizer(optimizer, path, args):
    '''Load model parameters from path'''
    
    checkpoint = torch.load(path, map_location='cpu')
    assert args.model_params == checkpoint['model_params'], \
        "Saved model params and current params are different"
    '''
    optimizer_state_dict = {k.partition('model.')[2]: v for k,v in \
                            checkpoint['optimizer_state_dict']}
    '''
    optimizer_state_dict = checkpoint['optimizer_state_dict']
    optimizer.load_state_dict(optimizer_state_dict)
    
    return optimizer




def reduce_learning_rate(optimizer):
    '''Reduce learning rate by a factor of 0.5'''

    optimizer.param_groups[0]['lr'] = optimizer.param_groups[0]['lr'] * 0.5
    
    return optimizer




def root_mean_squared_error(ytrue, ypred, weight):
    '''Return the weighted mean squared error'''
    
    return torch.sqrt(torch.mean(((ytrue - ypred) ** 2) * weight))
  



def prepare_training(args):
    '''Prepare paths and objects for training'''
    
    # Paths/directories
    args.savepath = f'{args.savedir}/{args.model_name}'
    if not os.path.exists(args.savepath):
        os.makedirs(args.savepath)
    args.recent_model = f'{args.savepath}/model_recent.pt' # Save model each epoch
    args.best_model = f'{args.savepath}/model_best.pt'     # Save model if error decreases
    args.log_file = f'{args.savepath}/log.csv'             # Write losses to file
    
    
    # Instantiate model 
    # Don't instantiate optimizer until model.to(device) is called
    model, args = build_model(args)
    
    
    # Model and training progress parameters
    if (not args.restart): # Train from scratch, don't restart
        
        # Write new log file
        with open(args.log_file, 'w') as logs:
            logs.write("epoch,train_loss,val_loss,best_val_loss,learning_rate,"\
                       "since_improved\n")
        
        # Progress parameters
        progress_params = {'epoch': 1, # Start from 1 to args.epoch + 1, not from 0
                           'best_val_loss': np.nan,
                           'since_improved': 0}
            
    else: # Restart training, don't train from scratch
        
        # Don't write new log files, assert that necessary files are in savepath
        error_msg = "Cannot restart training without "
        assert os.path.exists(args.log_file), error_msg + "progress file"
        assert os.path.exists(args.best_model), error_msg + "best_model"
        
        # Get progress parameters from log file
        logs = pd.read_csv(args.log_file, index_col=False, header=0)
        progress_params = {'epoch': logs['epoch'].values[-1],
                           'best_val_loss': logs['best_val_loss'].values[-1], 
                           'since_improved': logs['since_improved'].values[-1]}        

        # Load model from checkpoint
        model = load_model(model, args.best_model, args)
        
    return model, progress_params, args
     



def start_message(args):
    '''Description message at start of training'''
    
    if (args.verbose):
        print(f'INFO: MODEL_NAME = {args.model_name}')
        print(f'INFO: NUM_PARAMETERS = {args.num_parameters:.2e}')
        print(f'INFO: EPOCHS = {args.epochs}')
        print(f'INFO: NUM_NODES = {args.num_nodes}')
        print(f'INFO: GPUS_PER_NODE = {args.gpus_per_node}')
        print(f'INFO: NUM_GPUS = {args.world_size}')
        print(f'INFO: NUM_CPUS = {args.num_cpus}')
        print(f'INFO: NUM_WORKERS = {args.num_workers}')        
        print(f'INFO: SAVE_PATH = {args.savepath}')
        print()
        for key, value in args.params.items():
            print(f'PARAMS: {key}={value}')
        print()
        print('INFO: STARTING TRAINING')
        print()
        



def exit_message(epoch, args):
    '''Exit message at end of training'''
    
    if epoch >= args.epochs:
        print()
        print(f'INFO: Maximum training epochs ({args.epochs}) reached!')
        print('INFO: FINISHED TRAINING')
        
        return
    
    


def train_for_one_epoch(rank, trainloader, model, optimizer, epoch, args):
    '''Train model for one epoch'''
    
    _ = model.train()  # Set in training mode
    losses = []

    for i, (x, y, weight, mask) in enumerate(trainloader):

        x, y, weight, mask = x.to(rank), y.to(rank), weight.to(rank), mask.to(rank)
        optimizer.zero_grad()                          
        ypred = model(x, mask)[0]
        loss = root_mean_squared_error(y, ypred, weight) 
        _ = loss.backward()
        _ = optimizer.step()
        losses.append(loss.item())

    return np.mean(losses)




def validate(rank, valloader, model, args):
    '''Validate model on validation set'''
    
    _ = model.eval()  # Set in validation mode
    losses = []
    
    with torch.no_grad():
        
        for i, (x, y, weight, mask) in enumerate(valloader):
            
            x, y, weight, mask = x.to(rank), y.to(rank), weight.to(rank), mask.to(rank)
            ypred = model(x, mask)[0]
            loss = root_mean_squared_error(y, ypred, weight)
            losses.append(loss.item())
            
    return np.mean(losses)
            



def distributed_training(rank, world_size, args):
    '''Distributed training routine'''

    # Set up process groups for distributed training/inference
    _ = configure_processes(rank, args)
    
    
    # Prepare dataloader
    traindata = get_dataset(args.trainseqs, 
                            args.params['sample_method'], 
                            args)
    trainloader = get_distributed_dataloader(rank, 
                                             world_size,
                                             traindata,
                                             True,
                                             args)
    valdata = get_dataset(args.valseqs, 
                          'bin_inv', # Use bin_inv to reweight validataion data
                          args)
    valloader = get_distributed_dataloader(rank,
                                           world_size,
                                           valdata, 
                                           False,
                                           args)
    
    # Prepare paths and model objects for training
    model, progress_params, args = prepare_training(args)
    model = model.to(rank)  # Move model to GPU device in DDP
    if rank == 0:
        start_message(args)

    
    # Build optimizer 
    optimizer = build_optimizer(model, args)
    if args.restart:
        optimizer = load_optimizer(optimizer, args.best_model, args)
    

    # Wrap the model as a DistributedDataParallel object
    model = DistributedDataParallel(model, 
                                    device_ids=[rank],
                                    output_device=rank,
                                    find_unused_parameters=True)
    
    
    # Training epochs
    start_epoch = progress_params['epoch']  # start_epoch >0 if restarting
    stop_epoch = args.epochs + 1
    if rank == 0:
        _ = exit_message(start_epoch, args)
    
    # Training iteration
    for epoch in range(start_epoch, stop_epoch):
        
        _ = trainloader.sampler.set_epoch(epoch) # For distributed batches
        learning_rate = optimizer.param_groups[0]['lr']
        start_time = time.time()
        train_loss = train_for_one_epoch(rank, trainloader, model, optimizer,
                                                    epoch, args)
        val_loss = validate(rank, valloader, model, args)
        end_time = time.time()
        total_time = end_time - start_time 

        
        # Writing and priting, only if gpu is master
        if rank == 0:
            
            # Save model and update progress parameters
            _ = save_model(model, optimizer, args.recent_model, args)
            
            if not (val_loss >= progress_params['best_val_loss']):
                # Save model as best model, if validation loss has improved
                _ = save_model(model, optimizer, args.best_model, args)
                progress_params['best_val_loss'] = val_loss
                progress_params['since_improved'] = 0 # Reset
            else:
                progress_params['since_improved'] += 1 
                
            
            # Reduce learning rate if performance hasn't improved
            if (progress_params['since_improved']>= args.reduce_lr_patience) and \
                (progress_params['since_improved'] % args.reduce_lr_patience == 0):
                    optimizer = reduce_learning_rate(optimizer)
                
            # Log progress
            with open(args.log_file, 'a') as logs:
                logs.write(f"{epoch},{train_loss},{val_loss},"\
                           f"{progress_params['best_val_loss']},"\
                           f"{learning_rate},"\
                           f"{progress_params['since_improved']}\n")
                    
            if args.verbose == 1:
                print(f'PROGRESS: epoch={epoch}, '\
                      f'train_loss={train_loss:.4f}, '\
                      f'val_loss={val_loss:.4f}, '\
                      f"best_val_loss={progress_params['best_val_loss']:.4f}, "\
                      f"learning_rate={learning_rate:.2e}, "\
                      f"since_improved={progress_params['since_improved']}, "\
                      f"time={total_time:.1f}s")
            
            # Exit if validation hasn't improved 
            if progress_params['since_improved'] >= args.stop_patience:
                print(f"INFO: Exiting, as validation loss has not improved since "\
                      f"{progress_params['since_improved']} epochs")
                print("INFO: FINISHED TRAINING")
                return
                
    
    # Exit training
    if rank == 0:
        _ = exit_message(epoch, args)
    return




def main():
    '''Main routine'''
    
    # Parse command-line training arguments
    args = create_parser()                
    
    # Configure GPUs for processes
    world_size, args = configure_cuda(args)        
    
    # Distributed training raining routine
    _ = mp.spawn(distributed_training,
                 nprocs=args.world_size,
                 args=(args.world_size, args)
                 )




if __name__=='__main__':
    _ = main()