train.py

# template imports
import argparse
import collections
import torch
import numpy as np
# import data_loader.data_loaders as module_data
# import model.loss as module_loss
# import model.metric as module_metric
import models
# from parse_config import ConfigParser
from trainer import train
from utils import prepare_device, mkdir, get_logger 
from data_loader import get_data_loaders

# fix random seeds for reproducibility
SEED = 0
torch.manual_seed(SEED)
torch.backends.cudnn.deterministic = True
torch.backends.cudnn.benchmark = False
np.random.seed(SEED)

# nir added 22.6.21 from https://pytorch.org/docs/stable/notes/randomness.html
import random
random.seed(SEED)
# torch.use_deterministic_algorithms(True)   requires env var CUBLAS_WORKSPACE_CONFIG - see https://docs.nvidia.com/cuda/cublas/index.html#cublasApi_reproducibility



import logging



# external imports
import os
import time
from datetime import datetime
from argparse import ArgumentParser, ArgumentDefaultsHelpFormatter
import torch
import torch.nn as nn
# from torch.utils.tensorboard import SummaryWriter, summary # pytorch TB logger
# import optuna
import joblib
from copy import deepcopy
# import matplotlib.pyplot as plt
from glob import glob

# local imports
# import models
# from engines import create_trainer, create_evaluator

from torch.optim.lr_scheduler import CyclicLR, OneCycleLR, ReduceLROnPlateau, StepLR


# globals
task_time = datetime.now().strftime("%d%m%Y_%H%M%S")
output_dir = os.path.join('..','output','output_'+task_time)  # '/home/john/Desktop/git/scnn_ignite/output/output_'+task_time

# num_classes is good/bad
# categories_count is the number of uniq syscalls
# TODO: change categories_count to 46 for 60K_Max1000 or 48 without limit Max1000
configuration_data = {'num_classes': 2, 'categories_count': 46, 'embed_size': 300, 'hidden_size': 300}  # TODO: this dict is used outside of trains





def main(args):

    logger = get_logger(name='train')

    logger.info("######### starting main train.py #########")
    logger.info("######### args: ######### \n %s" %(args))

    # Define train, valid and test datasets
    logger.info("loading data")
    train_loader, val_loader, test_loader, future_loader = get_data_loaders(args.shuffled_epochs, args.train_procedure, args.dataset_path, args.samples_count, args.batch_size, args.max_sequence_length, args.time_weight)

    # Set the training device to GPU if available - if not set it to CPU
    # device = torch.cuda.current_device() if torch.cuda.is_available() else torch.device('cpu')
    # torch.backends.cudnn.benchmark = True if torch.cuda.is_available() else False  # optimization for fixed input size
    
    # prepare for (multi-device) GPU training
    # device, device_ids = prepare_device(1)
    # Set the training device to GPU if available - if not set it to CPU
    device = torch.cuda.current_device() if torch.cuda.is_available() else torch.device('cpu')
    torch.backends.cudnn.benchmark = True if torch.cuda.is_available() else False  # optimization for fixed input size

    # added an arg for BiGruLR or Transformer
    if args.architecture == 'Embed+AvgPool':
        logger.info("chosen architecture is Embed+AvgPool EmbedAvgPoolLR")
        # init model with config params
        model = models.EmbedAvgPoolLR(device=device,
                            batch_size=args.batch_size,
                            MaxSeqLen=args.max_sequence_length,
                            num_classes=configuration_data.get('num_classes'),
                            categories_count=configuration_data.get('categories_count'),
                            embed_size=args.embed_size,
                            hidden_size=args.hidden_size)
    elif args.architecture == 'GRU+LastPool':
        logger.info("chosen architecture is GRU+LastPool BiGruLastPoolLR")
        # init model with config params
        model = models.BiGruLastPoolLR(device=device,
                            batch_size=args.batch_size,
                            MaxSeqLen=args.max_sequence_length,
                            num_classes=configuration_data.get('num_classes'),
                            categories_count=configuration_data.get('categories_count'),
                            embed_size=args.embed_size,
                            hidden_size=args.hidden_size)
    elif args.architecture == 'LSTM+LastPool':
        logger.info("chosen architecture is LSTM+LastPool BiLstmLastPoolLR")
        # init model with config params
        model = models.BiLstmLastPoolLR(device=device,
                            batch_size=args.batch_size,
                            MaxSeqLen=args.max_sequence_length,
                            num_classes=configuration_data.get('num_classes'),
                            categories_count=configuration_data.get('categories_count'),
                            embed_size=args.embed_size,
                            hidden_size=args.hidden_size)
    elif args.architecture == 'LSTM+AvgPool':
        logger.info("chosen architecture is LSTM+AvgPool BiGruAvgPoolLR")
        # init model with config params
        model = models.BiGruAvgPoolLR(device=device,
                            batch_size=args.batch_size,
                            MaxSeqLen=args.max_sequence_length,
                            num_classes=configuration_data.get('num_classes'),
                            categories_count=configuration_data.get('categories_count'),
                            embed_size=args.embed_size,
                            hidden_size=args.hidden_size)
    elif args.architecture == 'GRU+AvgPool':
        logger.info("chosen architecture is GRU+AvgPool BiGruAvgPoolLR")
        # init model with config params
        model = models.BiGruAvgPoolLR(device=device,
                            batch_size=args.batch_size,
                            MaxSeqLen=args.max_sequence_length,
                            num_classes=configuration_data.get('num_classes'),
                            categories_count=configuration_data.get('categories_count'),
                            embed_size=args.embed_size,
                            hidden_size=args.hidden_size)
    elif args.architecture == 'CNN':
        logger.info("chosen architecture is CNN CnnLR")
        # init model with config params
        model = models.CnnLR(device=device,
                            batch_size=args.batch_size,
                            MaxSeqLen=args.max_sequence_length,
                            num_classes=configuration_data.get('num_classes'),
                            categories_count=configuration_data.get('categories_count'),
                            embed_size=args.embed_size,
                            hidden_size=args.hidden_size)
    elif args.architecture == 'LSTM+MaxPool':
        logger.info("chosen architecture is LSTM+MaxPool BiLstmLR")
        # init model with config params
        model = models.BiLstmLR(device=device,
                            batch_size=args.batch_size,
                            MaxSeqLen=args.max_sequence_length,
                            num_classes=configuration_data.get('num_classes'),
                            categories_count=configuration_data.get('categories_count'),
                            embed_size=args.embed_size,
                            hidden_size=args.hidden_size)
    elif args.architecture == 'LSTM+Attention':
        logger.info("chosen architecture is LSTM+Attention BiLstmAttnLR")
        # init model with config params
        model = models.BiLstmAttnLR(device=device,
                            batch_size=args.batch_size,
                            MaxSeqLen=args.max_sequence_length,
                            num_classes=configuration_data.get('num_classes'),
                            categories_count=configuration_data.get('categories_count'),
                            embed_size=args.embed_size,
                            hidden_size=args.hidden_size)                            
    elif args.architecture == 'GRU+Transformer':
        logger.info("chosen architecture is GRU+Transformer BiGruGptCls")
        # init model with config params
        model = models.BiGruGptCls(device=device,
                                   batch_size=args.batch_size,
                                   MaxSeqLen=args.max_sequence_length,
                                   categories_count=configuration_data.get('categories_count'),
                                   embedding_size=args.embed_size,
                                   hidden_size=args.hidden_size,
                                   classes_count=configuration_data.get('num_classes'),
                                   normalize=False)
    elif args.architecture == 'Transformer':
        logger.info("chosen architecture is Transformer GPT2")
        # init model with config params
        model = models.GPT2(device=device,
                                   batch_size=args.batch_size,
                                   MaxSeqLen=args.max_sequence_length,
                                   categories_count=configuration_data.get('categories_count'),
                                   embedding_size=args.embed_size,
                                   hidden_size=args.hidden_size,
                                   classes_count=configuration_data.get('num_classes'),
                                   normalize=False)
    elif args.architecture == 'GRU+Attention':
        logger.info("chosen architecture is GRU+Attention (BiGruAttnLR)")
        # init model with config params
        model = models.BiGruAttnLR(device=device,
                               batch_size=args.batch_size,
                               MaxSeqLen=args.max_sequence_length,
                               num_classes=configuration_data.get('num_classes'),
                               categories_count=configuration_data.get('categories_count'),
                               embed_size=args.embed_size,
                               hidden_size=args.hidden_size)
    else:
        logger.info("defualt architecture is GRU+MaxPool (BiGruLR)")
        # init model with config params
        model = models.BiGruLR(device=device,
                               batch_size=args.batch_size,
                               MaxSeqLen=args.max_sequence_length,
                               num_classes=configuration_data.get('num_classes'),
                               categories_count=configuration_data.get('categories_count'),
                               embed_size=args.embed_size,
                               hidden_size=args.hidden_size)
    
    # get function handles of loss
    criterion = nn.NLLLoss()

    # init training objects
    # trainable_params = filter(lambda p: p.requires_grad, model.parameters())
    trainable_params = [p for p in model.parameters() if p.requires_grad]
    optimizer = torch.optim.Adam(trainable_params, lr=args.learning_rate, weight_decay=args.weight_decay)
    if args.lr_scheduler == 'ReduceLROnPlateau':
        lr_scheduler = torch.optim.lr_scheduler.ReduceLROnPlateau(optimizer,
                                                                mode='max',  # TODO: change by LR_scheduler_metric !
                                                                factor=0.1,  # 0.1,
                                                                patience=2,  # 2,
                                                                threshold=1e-4,  # 1e-4,
                                                                min_lr=1e-15,
                                                                verbose=True)
    elif args.lr_scheduler == 'StepLR':
        lr_scheduler = torch.optim.lr_scheduler.StepLR(optimizer, step_size=1, gamma=0.5)                                                              
    elif args.lr_scheduler == 'CyclicLR':
        lr_scheduler = torch.optim.lr_scheduler.CyclicLR(optimizer, 
                                                         base_lr=1e-8, 
                                                         max_lr=args.learning_rate, 
                                                         cycle_momentum=False, 
                                                         step_size_up = len(train_loader)/2,
                                                         mode="triangular2")
    elif args.lr_scheduler == 'OneCycleLR':
        lr_scheduler = torch.optim.lr_scheduler.OneCycleLR(optimizer, 
                                                           max_lr=args.learning_rate, 
                                                           total_steps = 30*len(train_loader),
                                                           cycle_momentum=False, 
                                                           final_div_factor=args.learning_rate/1e-8)

    
    # prepare for (multi-device) GPU training
    model = model.to(device)
    # if len(device_ids) > 1:
        # model = torch.nn.DataParallel(model, device_ids=device_ids)


    
    train(args, model, criterion, optimizer, lr_scheduler, device, train_loader, val_loader, test_loader, future_loader)

    logger.info("####### finished [main] train.py ! ####### ")



if __name__ == '__main__':
    # init arg parse
    parser = argparse.ArgumentParser(formatter_class=ArgumentDefaultsHelpFormatter)
    
    # template args
    # parser.add_argument('-r', '--resume', default=None, type=str,
    #                   help='path to latest checkpoint (default: None)')
    # parser.add_argument('-d', '--device', default=None, type=str,
    #                   help='indices of GPUs to enable (default: all)')

    # more args 
    parser.add_argument("-e", "--epochs", required=False, type=int, default=2, help="number of epochs to train")
    parser.add_argument("-sc", "--samples_count", required=False, type=int, default=0,
                        help="amount of syscall sequences for training")  # 1000
    parser.add_argument("-dp", "--dataset_path", required=True, type=str,
                        help="dataset directory path for scnn (with train, valid and test folders inside)")
    parser.add_argument("-bs", "--batch_size", required=False, type=int, default=30,
                        help="amount of syscalls for each batch")
    parser.add_argument("-msl", "--max_sequence_length", required=False, type=int, default=1000,  # TODO: change maxlen to the syscall_mapping1_fix_Max1000.pickle
                        help="maximum syscall sequence length")
    parser.add_argument("-dt", "--dataset_type", required=False, type=str, default="tensor_csv",  # TODO: delete!
                        help="type of the dataset to load for scnn (binee, random, cuckoo, seq)")
    parser.add_argument("-esp", "--early_stopping_patience", required=False, type=int, default=-1,
                        help="patience for the early stopping of the training phase, default is epoch_count (won't stop)")
    parser.add_argument("-mf", "--max_fpr", required=False, type=float, default=0.02,
                        help="fpr threshold for early stopping of the training phase (float, not percentage)")
    parser.add_argument("-esm", "--early_stopping_metric", required=False, type=str, default='accuracy',
                        help="metric to consider for the early stopping of the training phase")
    parser.add_argument('-sched', "--lr_scheduler", type=str, default='ReduceLROnPlateau',
                        choices=['StepLR', 'ReduceLROnPlateau', 'CyclicLR', 'OneCycleLR'],
                        help=("--lr_scheduler = 'StepLR', 'ReduceLROnPlateau', 'CyclicLR', 'OneCycleLR'"))
    parser.add_argument("-lrsm", "--LR_scheduler_metric", required=False, type=str, default='ROC_AUC',
                        help="metric to consider for the early stopping of the training phase")
    parser.add_argument("-wd", "--weight_decay", required=False, type=float, default=0.01,
                        help="learning rate of the model")
    parser.add_argument("-lr", "--learning_rate", required=False, type=float, default=1e-5,
                        help="learning rate of the model")
    parser.add_argument('--warmup_iterations', type=int, default=5000,
                        help='Number of iteration for warmup period (until reaching base learning rate)')
    parser.add_argument('--log_interval', type=int, default=100,
                        help='how many batches to wait before logging training status')
    parser.add_argument('-i', '--input_checkpoint', type=str, default='',
                        help='Loading model, optimizer, and LR-Scheduler from checkpoint.')
    parser.add_argument("--output_dir", type=str, default=output_dir,
                        help="output directory for saving model checkpoints and tensorboard_logs")
    parser.add_argument("--momentum", type=float, default=0.9,
                        help="momentum for optimizer")

    parser.add_argument("-opt", "--optimize_params", action='store_true', required=False, default=False,
                        help="optimize-params with optuna")
    parser.add_argument("-prun", "--opt_pruning", action='store_true', required=False, default=False,
                        help="enable trials pruning with optuna")
    parser.add_argument('-trials', '--opt_trials', type=int, default=None,
                        help='number of trials to optimize-params with optuna')
    parser.add_argument('-timeout', '--opt_timeout', type=float, default=None,
                        help='number of seconds to optimize-params with optuna; None=until ^C or signal.')

    parser.add_argument('-emb', '--embed_size', type=int, default=configuration_data['embed_size'],
                        help='NN embed_size')
    parser.add_argument('-hid', '--hidden_size', type=int, default=configuration_data['hidden_size'],
                        help='NN hidden_size')

    parser.add_argument('-tw', '--time_weight', type=float, default=0,
                        help='time_weight for the samples weight')

    # TODO: lookout from the similiar args - train/trains
    parser.add_argument("-trains", "--allegro_trains", action='store_true', required=False, default=False,
                        help="ATTENTION: this uses allegro-trains server (defaults to upload code to demo server!! unless local-server configured)")

    parser.add_argument('-t', "--train_procedure", type=str, default='train',
                        choices=['train', 'retrain', 'double_retrain', 'all'],
                        help=("train: train on train-set, valid on valid-set, test on test-set. Do test also on future-set." + \
                             "retrain: train on (train-set+valid-set), valid on test-set. Do test on future-set." + \
                             "double-retrain: train on (train-set+valid-set+test-set). Do test on future-set." + \
                              "all: running 3 executions - train + retrain + double_retrain."))

    parser.add_argument('-checkpoint', "--save_checkpoints", type=str, default='none',
                        choices=['none', 'all', 'lr'],
                        help=("none: don't save any state each epoch." + \
                             "all: save model+optimizer+LR_scheduler checkpoints (each epoch store 12MB file to disk)." + \
                             "lr: save LR_scheduler checkpoints (each epoch store 12MB file to disk)"))

    parser.add_argument("-shuffle", "--shuffled_epochs", action='store_true', required=False, default=False,
                        help="set DataLoader argument shuffle=True, used for shuffled_epochs in which the interior order of samples inside the set is random and not time sorted.")

    parser.add_argument('-arch', "--architecture", type=str, default='GRU',
                    choices=['Embed+AvgPool', 'GRU+LastPool', 'LSTM+LastPool', 'GRU+AvgPool', 'LSTM+AvgPool', 'CNN', 'LSTM+MaxPool', 'LSTM+Attention', 'GRU+MaxPool', 'GRU+Attention', 'Transformer', 'GRU+Transformer'],
                    help=("model architecture:\n 'CNN'.\n 'LSTM' or 'GRU' with '+Attention' or '+(Max/Avg/Max)Pool'.\n 'Transformer' and 'GRU+Transformer'."))

    args = parser.parse_args()

    # init stuf
    # set_max_fpr(args.max_fpr)
    output_dir = args.output_dir
    mkdir(args.output_dir)

    # if allergro-Trains:
    if args.allegro_trains:
        from trains import Task
        task = Task.init(project_name='SCNN with TRAINS, Ignite and TensorBoard',  # Task.create
                         task_name=('Train SCNN with tensor_csv dataset: ' + task_time),
                         output_uri=output_dir)  # output_'+task_time
        configuration_data = task.connect_configuration(configuration_data)

    start_time = time.time()
    logger = get_logger(name='train', dir=args.output_dir)
    if args.optimize_params:
        optimized_study = optimize_study(args)
        joblib.dump(optimized_study, os.path.join(args.output_dir, 'optimized_study.pkl'))
    else:                                                                       
        # in case of args.train == all -> run the args 3 times: train + retrain + double_retrain
        if args.train_procedure == 'all':
            logger.info("\n\n\n ### running all training-procedures: ###\n\n\n")
            for procedure in ['train', 'retrain', 'double_retrain']:
                procedure_args = deepcopy(args)
                procedure_args.train_procedure = procedure                
                # force lr_checkpoints of the retrain-procedure for double-retrain lr-loading
                if procedure == 'retrain' and procedure_args.save_checkpoints == 'none':
                    procedure_args.save_checkpoints = 'lr'
                # create procedure dir
                procedure_args.output_dir = os.path.join(args.output_dir, procedure)
                mkdir(procedure_args.output_dir)
                # run procedurrre
                logger.info("\n\n\n ### starting train-procedure: %s ###\n\n\n" %(procedure))
                procedure_start_time = time.time()
                main(procedure_args)
                logger.info("\n\n\n ### train-procedure: %s finished in %d seconds ###\n\n\n" % (procedure, time.time() - procedure_start_time))
        else:
            logger.info("\n\n\n ### running "+args.train_procedure+" training-procedure: ###\n\n\n")
            # create procedure dir
            args.output_dir = os.path.join(args.output_dir, args.train_procedure)
            mkdir(args.output_dir)
            main(args)
    logger.info("--- that took %s seconds ---" % (time.time() - start_time))