train.py

import os
import json
import math
import torch
import random
import logging
import numpy as np
from collections import OrderedDict
from tensorboardX import SummaryWriter
from torch.utils.data.distributed import DistributedSampler
from torch.utils.data import DataLoader, RandomSampler, SequentialSampler, TensorDataset
from transformers import (
    WEIGHTS_NAME,
    AdamW,
    get_linear_schedule_with_warmup
)
from processors import processor_dict
from models import load_model, save_model
from evaluate import evaluate_func_dict, eval_checkpoint
from torch.nn import CrossEntropyLoss

def set_seed(seed):
    random.seed(seed)
    np.random.seed(seed)
    torch.manual_seed(seed)
    torch.cuda.manual_seed_all(seed)
    

def get_optimizer(model, weight_decay, learning_rate, adam_epsilon):
    no_decay = ["bias", "LayerNorm.weight"]
    optimizer_grouped_parameters = [
        {"params": [p for n, p in model.named_parameters() if not any(nd in n for nd in no_decay)],
        "weight_decay": weight_decay},
        {"params": [p for n, p in model.named_parameters() if any(nd in n for nd in no_decay)], "weight_decay": 0.0}
    ]
    optimizer = AdamW(optimizer_grouped_parameters, lr=learning_rate, eps=adam_epsilon)
    return optimizer



def score_comp(score_a, score_b, metric_comp):
    if metric_comp == "larger":
        return score_a > score_b
    elif metric_comp == "smaller":
        return score_a < score_b
    else:
        raise NotImplementedError()


def evaluation_step(
            args, 
            model, tokenizer, processor, 
            optimizer, scheduler,
            dev_file_infos, test_file_infos, 
            eval_func, 
            global_step, 
            best_dev_score,
            test_score_on_best_dev,
            tb_writer=None,
            tf_board_header="",
            logger=None
        ):
    if logger == None:
        logger = logging.getLogger(__name__)
    dev_dataset_results = eval_checkpoint(
                    args, 
                    model, tokenizer, processor,
                    global_step=global_step,
                    file_infos=dev_file_infos, 
                    eval_func=eval_func, 
                    tb_writer=None,
                    tf_board_header=tf_board_header,
                    mode="dev",
                    logger=logger
                )
    logger.info("Dev Evauation : {}".format(json.dumps(dev_dataset_results, indent=4)))
    tb_writer.add_scalar(
                    "{}_avg_dev".format(tf_board_header), 
                    dev_dataset_results["avg_metric"],
                    global_step
                )
    if  (score_comp(dev_dataset_results["avg_metric"], best_dev_score["avg_metric"], args.dev_metric_comp)):
        best_dev_score = dev_dataset_results
        save_model(
            args.output_dir, 
            args, 0, 
            model, tokenizer, 
            optimizer, scheduler, 
            logger, prefix="best_dev"
        )
        logger.info("New BEST-DEV score found")
        open(os.path.join(args.output_dir, "best_dev_score.json"), "w").write(json.dumps(best_dev_score, indent=4))
        
        if args.evaluate_test_on_best_dev:
            test_dataset_results = eval_checkpoint(
                        args, 
                        model, tokenizer, processor,
                        global_step=global_step,
                        file_infos=test_file_infos, 
                        eval_func=eval_func, 
                        tb_writer=None,
                        tf_board_header=tf_board_header,
                        mode="test",
                        logger=logger
                    )
            logger.info("Test Evauation : {}".format(json.dumps(test_dataset_results, indent=4)))
            test_score_on_best_dev = test_dataset_results
            open(os.path.join(args.output_dir, "test_score_on_best_dev.json"), "w").write(json.dumps(test_score_on_best_dev, indent=4))
    
    return best_dev_score, test_score_on_best_dev


def training_loop(
        args, 
        train_datasets, 
        model, tokenizer, processor,
        optimizer, scheduler, 
        dev_dataset_results,
        eval_func=None,
        dev_file_infos=None,
        test_file_infos=None,
        tf_board_header="",
        logger=None
    ):
    if logger is None:
        logger = logging.getLogger(__name__)
    
    tb_writer = None
    if args.local_rank in [-1, 0]:
        # tensorboard setup
        logger.info("Tensorboard summary creation.")
        tb_path = os.path.join(args.output_dir, "tf_board")
        if not os.path.exists(tb_path):
            os.makedirs(tb_path)
        tb_writer = SummaryWriter(tb_path)

    # At first evaluate the result
    test_score_on_best_dev = None,
    if ( args.local_rank == -1 and args.evaluate_during_training):
        dev_dataset_results, test_score_on_best_dev = evaluation_step(
                        args, 
                        model, tokenizer, processor, 
                        optimizer, scheduler,
                        dev_file_infos, test_file_infos, 
                        eval_func, 
                        global_step=0, 
                        best_dev_score=dev_dataset_results,
                        test_score_on_best_dev=test_score_on_best_dev,
                        tb_writer=tb_writer,
                        tf_board_header=tf_board_header,
                        logger=logger
                    )

    model.zero_grad()

    # batch size setup
    args.train_batch_size = args.per_gpu_train_batch_size * max(1, args.n_gpu)

    logger.info("Iteratior Sampling rate calculation.")
    
    # iteerator sampling rate
    (dataset_len, 
     prob_index, 
     prob_list) = iterator_selection_prob(
         args.sampling_penalty, 
         train_datasets, 
         logger
    )

    logger.info("Processing data loader.")
    
    # Data loader creation
    train_data_loader = []
    total_num_of_samples = 0
    for k in prob_index:
        dataset = train_datasets[k]
        total_num_of_samples += len(dataset)
        train_sampler = RandomSampler(dataset) if args.local_rank == -1 else DistributedSampler(dataset)
        data_loader = DataLoader(
                dataset, sampler=train_sampler, batch_size=args.per_gpu_train_batch_size
            )
        train_data_loader.append((k, data_loader))
    
    logger.info("Processing initial iterator.")
    # Initial Iterator creation
    train_iterators = []
    for i in range(len(train_data_loader)):
        assert train_data_loader[i][0] == prob_index[i]
        train_iterators.append(iter(train_data_loader[i][1]))
    tot_num_of_iterator = len(train_iterators)
    num_of_batch_trained = [ 0 for i in range(tot_num_of_iterator) ]
    
    # Number of steps calc
    if args.max_steps > 0 :
        total_training_steps = args.max_steps
    else:
        total_training_steps = int((total_num_of_samples * args.num_train_epochs) // (args.per_gpu_train_batch_size * args.gradient_accumulation_steps * torch.cuda.device_count()))
    
    logger.info("Declare Optimizer and Scheduler if needed.")
    # Declaring optimizer and scheduler
    optimizer = get_optimizer(
                    model, 
                    args.weight_decay, 
                    args.learning_rate, 
                    args.adam_epsilon
                ) if optimizer is None else optimizer
    if args.fp16:
        scaler = torch.cuda.amp.GradScaler() 
    if args.warmup_steps == -1:
        assert args.warmup_percentage >=0 and args.warmup_percentage < 20
        args.warmup_steps = int(max(0, (total_training_steps * args.warmup_percentage)//100))
    scheduler = get_linear_schedule_with_warmup(
                    optimizer, 
                    num_warmup_steps=args.warmup_steps, 
                    num_training_steps=total_training_steps
                ) if  scheduler is None else scheduler
    
    
    logger.info("Parallalize model and data.")
    # Parallalize model and dataset
    if args.n_gpu > 1:
        model = torch.nn.DataParallel(model)

    if args.local_rank != -1:
        model = torch.nn.parallel.DistributedDataParallel(
            model, device_ids=[args.local_rank], output_device=args.local_rank, find_unused_parameters=True
        )
    if args.local_rank in [-1, 0]:
        # Training log
        logger.info("***** Running training *****")
        logger.info("  Num examples = {}".format(total_num_of_samples))
        logger.info("  Num Epochs = {}".format(args.num_train_epochs))
        logger.info("  Instantaneous batch size per GPU = {}".format(args.per_gpu_train_batch_size))
        logger.info(
            "  Total train batch size (w. parallel, distributed & accumulation) = {}".format(
                    args.train_batch_size
                    * args.gradient_accumulation_steps
                    * (torch.distributed.get_world_size() if args.local_rank != -1 else 1)
            )
        )
        logger.info("  Gradient Accumulation steps = {}".format(args.gradient_accumulation_steps))
        logger.info("  Total optimization steps = {} (forward steps : {})".format(
                    total_training_steps, 
                    total_training_steps*args.gradient_accumulation_steps
                        )
                    )
        logger.info("  Total warmup steps = {}".format(args.warmup_steps))

    # Resume training
    saved_global_step=-1
    if os.path.exists(args.model_name_or_path) and args.resume_training:
        saved_global_step = int(args.model_name_or_path.split("-")[-1].split("/")[0])
        logger.info("  Continuing training from checkpoint, will skip to saved global_step : {}".format(saved_global_step))

    
    # Training var declaration
    tr_loss, logging_loss = 0.0, 0.0
    steps, global_step = 0, 0
    isUpdated = 0
    
    # Precalculate iterator id for each steps
    iterator_ids = []
    total_number_of_batch = int(total_training_steps*args.gradient_accumulation_steps)
    for _i in range(total_number_of_batch):
        if _i %  args.gradient_accumulation_steps == 0:
            _iterator_id = np.random.choice(range(tot_num_of_iterator), p=prob_list)
        iterator_ids.append(
            _iterator_id
        )

    # Added here for reproductibility
    set_seed(args.seed) 

    for iterator_ids_index in range(total_number_of_batch):
        iterator_id = iterator_ids[iterator_ids_index]
        model.train()
        try:
            batch = train_iterators[iterator_id].__next__()
        except StopIteration:
            train_iterators[iterator_id] = iter(train_data_loader[iterator_id][1])
            batch = train_iterators[iterator_id].__next__()
        num_of_batch_trained[ iterator_id ] += 1

        # No need to train if current step is less than saved_global_step
        if saved_global_step > global_step:
            steps = steps + 1
            if steps == args.gradient_accumulation_steps:
                steps = 0
                global_step = global_step + 1
            continue
        
        # Prepare batch
        batch = tuple(t.to(args.device) for t in batch)
        inputs = {
                    "input_ids": batch[0],
                    "attention_mask": batch[1],
                    "labels": batch[3],
                    "conf_penalty": args.conf_penalty if args.conf_penalty == 1 else None,
                    "marginal_entropy": args.marginal_entropy if args.marginal_entropy == 1 else None
            }
        if args.model_type != "distilbert":
            # XLM and RoBERTa don't use segment_ids
            inputs["token_type_ids"] = batch[2] if args.model_type in ["bert", "xlnet"] else None  

        # Inference
        loss, _ = model(**inputs)
        if args.n_gpu > 1:
            loss = loss.mean()  
        if args.gradient_accumulation_steps > 1:
            loss = loss / args.gradient_accumulation_steps
        # fp16 training
        if args.fp16:
            scaler.scale(loss).backward()
        else:
            loss.backward()
        tr_loss += loss.item()

        if (iterator_ids_index+1) % args.gradient_accumulation_steps == 0:
            # fp16 training
            if args.fp16:
                scaler.unscale_(optimizer)
                
            torch.nn.utils.clip_grad_norm_(model.parameters(), args.max_grad_norm)

            # training step
            if args.fp16:
                scaler.step(optimizer)
                scaler.update()
            else:
                optimizer.step()
            scheduler.step()
            model.zero_grad()
            global_step += 1

            if args.local_rank in [-1, 0] and args.logging_steps > 0 and global_step % args.logging_steps == 0:
                # Log metrics
                current_loss = (tr_loss - logging_loss) / args.logging_steps
                tb_writer.add_scalar("lr", scheduler.get_lr()[0], global_step)
                tb_writer.add_scalar("loss", current_loss , global_step)
                logger.info("[GLOBAL_STEP] : {}/{}, [LOSS] :: {}".format(global_step, total_training_steps, current_loss))
                logging_loss = tr_loss

            # save checkpoint
            if args.local_rank in [-1, 0] and args.save_steps > 0 and global_step % args.save_steps == 0:
                save_model(args.output_dir, args, global_step, model, tokenizer, optimizer, scheduler, logger)
            
            if ( args.local_rank in [-1, 0] and \
                args.evaluate_during_training and \
                    args.eval_steps > 0 and \
                        global_step % args.eval_steps == 0):
                
                dev_dataset_results, test_score_on_best_dev = evaluation_step(
                                args, 
                                model, tokenizer, processor, 
                                optimizer, scheduler,
                                dev_file_infos, test_file_infos, 
                                eval_func, 
                                global_step=global_step, 
                                best_dev_score=dev_dataset_results,
                                test_score_on_best_dev=test_score_on_best_dev,
                                tb_writer=tb_writer,
                                tf_board_header=tf_board_header,
                                logger=logger
                            )
                for enum_idx, (dataset_name, dataloader_obj) in enumerate(train_data_loader):
                    logger.info("TRAIN DATASET : {} : {} batch trained.".format(dataset_name, num_of_batch_trained[ enum_idx ]))
                logger.info("Best-Dev : {}".format(json.dumps(dev_dataset_results, indent=4)))
                if args.evaluate_test_on_best_dev:
                    logger.info("Test score on Best-Dev : {}".format(json.dumps(test_score_on_best_dev, indent=4)))

    if args.local_rank in [-1, 0]:
        tb_writer.close()

    return global_step, tr_loss / global_step
    

def baseline_training(args, logger=None):
    if logger is None:
        logger = logging.getLogger(__name__)
    if args.task_name not in processor_dict:
        raise ValueError("Task not found in processor pool: {}".format((args.task_name)))
    if args.task_name not in evaluate_func_dict:
        raise ValueError("Task not found eval function pool: {}".format((args.task_name)))
    processor = processor_dict[args.task_name](
                    src_lang=args.src_lang, 
                    dev_lang=args.dev_lang, 
                    tgt_lang=args.tgt_lang, 
                    seed=args.seed,
                    percentage=args.train_data_percentage,
                    n_shot=None,
                    pad_token_label_id=CrossEntropyLoss().ignore_index
                )
    args.output_mode = processor.output_mode

    label_list = processor.get_labels()
    num_labels = len(label_list)

    if args.local_rank not in [-1, 0]:
        torch.distributed.barrier()
    
    (config, 
     tokenizer, 
     model,
     optimizer, 
     scheduler) = load_model(
                task_name=args.task_name, 
                config_name=args.config_name, 
                tokenizer_name=args.tokenizer_name, 
                model_name_or_path=args.model_name_or_path, 
                num_labels=num_labels, 
                model_type=args.model_type, 
                logger=logger,
                do_lower_case=args.do_lower_case, 
                cache_dir=args.cache_dir,
                is_fp16=args.fp16
        )
    logger.info("Successfully loaded the config : {}, tokenizer : {}, model : {}".format(
                args.config_name, args.tokenizer_name, args.model_name_or_path))
    
    if args.local_rank == 0:
        torch.distributed.barrier()

    evaluate_func = evaluate_func_dict[args.task_name]
    model.to(args.device)
    train_datasets = {}
    for data_file_info in args.train:
        if data_file_info.split(";")[-1] in args.src_lang.split(";"):
            logger.info("Process rank : {} , Loading : {}".format(args.local_rank, data_file_info))
            train_dataset = processor.load_and_cache_examples(args, tokenizer, "train", data_file_info, logger)
            train_datasets[data_file_info] = train_dataset
    
    dev_file_infos, test_file_infos = [], []
    for data_file_info in args.dev:
        if data_file_info.split(";")[-1] in args.dev_lang.split(";"):
            dev_file_infos.append(data_file_info)
    for data_file_info in args.test:
        if data_file_info.split(";")[-1] in args.tgt_lang.split(";"):
            test_file_infos.append(data_file_info)
    
    global_step, tr_loss = training_loop(
            args, train_datasets, 
            model, tokenizer, processor,
            optimizer, scheduler,
            dev_dataset_results={"avg_metric":0},
            eval_func=evaluate_func,
            dev_file_infos=dev_file_infos if args.evaluate_during_training else None,
            test_file_infos=test_file_infos if (args.evaluate_during_training and args.evaluate_test_on_best_dev) else None,
            tf_board_header="{}_baseline".format(args.task_name),
            logger=logger
        )
    logger.info(" global_step = %s, average loss = %s", global_step, tr_loss)