train_multitask.py

import argparse
import math
import os
from datetime import datetime

import numpy as np
import torch
from torch.utils.data import DataLoader, Subset
from datasets import load_metric
from transformers import (AdamW, AutoConfig, AutoTokenizer, SchedulerType,
                          get_scheduler)

from metrics import task_to_metric_name
from data_loader import task_to_benchmark, task_to_collator, task_to_load_fns, ws_task_to_load_fns
from data_loader.multitask_dataset import (MultitaskBatchSampler,
                                           MultitaskCollator, MultitaskDataset)
from models.modeling_multi_bert import MultitaskBertForClassification
from utils import add_result_to_csv, get_logger, setup_logging
from trainer import MultitaskTrainer

def setup_logging_logic():
    # set up log file
    log_dir = "./saved/logs/{}".format(datetime.now().strftime(r'%m%d_%H%M%S'))
    if not os.path.exists(log_dir):
        os.mkdir(log_dir)
    setup_logging(log_dir)

def load_ws_task_data(task_name, model_name_or_path, pad_to_max_length, max_length, batch_size, lf_idxes, ws_method = "none", ws_params = {},
                    downsample_frac=1.0):
    config = AutoConfig.from_pretrained(model_name_or_path)
    tokenizer = AutoTokenizer.from_pretrained(model_name_or_path, use_fast=True)

    train_datasets, valid_dataset, test_dataset, num_lfs, num_labels = ws_task_to_load_fns[task_name](
            task_name=task_name, tokenizer = tokenizer,
            pad_to_max_length=pad_to_max_length, max_length=max_length, ws_method=ws_method, ws_params=ws_params)
    assert min(lf_idxes) >= 0 and max(lf_idxes) <= num_lfs

    rng = np.random.default_rng(int.from_bytes(task_name.encode("utf-8"), "little"))
    if 0 < downsample_frac < 1.0:
        for i, train_dataset in enumerate(train_datasets):
            tmp_len = len(train_dataset)
            if tmp_len > 100: # skip ws dataset with size less than 100
                downsample_len = int(tmp_len*downsample_frac) 
                indices = rng.choice(tmp_len, downsample_len, replace=False)
                train_datasets[i] = Subset(train_dataset, indices=indices)
    print("Data set lengths: " + " ".join([str(len(dataset)) for dataset in train_datasets]))

    train_dataloaders = [
        DataLoader(train_dataset, collate_fn=task_to_collator[task_name], batch_size=batch_size) for train_dataset in train_datasets
        ]
    valid_dataloader = DataLoader(valid_dataset, collate_fn=task_to_collator[task_name], batch_size=batch_size)
    test_dataloader = DataLoader(test_dataset, collate_fn=task_to_collator[task_name], batch_size=batch_size)

    task_to_train_datasets = dict([
        ("{}_{}".format(task_name, i), train_datasets[i]) for i in lf_idxes
        ])
    task_to_train_dataloaders = dict([
        ("{}_{}".format(task_name, i), train_dataloaders[i]) for i in lf_idxes
        ])
    task_to_valid_dataloaders = dict([
        ("{}_{}".format(task_name, i), valid_dataloader) for i in [0] # + lf_idxes
        ])
    task_to_test_dataloaders = dict([
        ("{}_{}".format(task_name, i), test_dataloader) for i in [0] # + lf_idxes
        ])
    task_to_num_labels = dict([
        ("{}_{}".format(task_name, i), num_labels) for i in [0] + lf_idxes
        ])
    
    task_to_metrics = dict([
        ("{}_{}".format(task_name, i), load_metric(task_to_metric_name[task_name], task_name)) for i in [0] + lf_idxes
        ])

    multitask_train_dataset = MultitaskDataset(task_to_train_datasets)
    multitask_train_sampler = MultitaskBatchSampler(task_to_train_datasets, batch_size)
    multitask_train_collator = MultitaskCollator(task_to_collator)
    multitask_train_dataloader = DataLoader(
        multitask_train_dataset,
        batch_sampler=multitask_train_sampler,
        collate_fn=multitask_train_collator.collator_fn,
    )
    return multitask_train_dataloader, task_to_train_dataloaders, task_to_valid_dataloaders, task_to_test_dataloaders, \
        task_to_num_labels, task_to_metrics, config, tokenizer 

def load_task_data(task_names, model_name_or_path, pad_to_max_length, max_length, batch_size, 
                    downsample_frac=1.0):
    config = AutoConfig.from_pretrained(model_name_or_path)
    tokenizer = AutoTokenizer.from_pretrained(model_name_or_path, use_fast=True)

    task_to_train_datasets = {}
    task_to_train_dataloaders = {}
    task_to_valid_dataloaders = {}
    task_to_test_dataloaders = {}
    task_to_num_labels = {}

    for task_name in task_names:
        train_dataset, valid_dataset, test_dataset, num_labels = task_to_load_fns[task_name](
            task_name=task_name, benchmark_name=task_to_benchmark[task_name],
            config = config, tokenizer = tokenizer,
            pad_to_max_length=pad_to_max_length, max_length=max_length)

        if 0 < downsample_frac < 1.0:
            tmp_len = len(train_dataset)
            if tmp_len > 100: # skip ws dataset with size less than 100
                rng = np.random.default_rng(int.from_bytes(task_name.encode("utf-8"), "little"))
                downsample_len = int(tmp_len*downsample_frac) 
                indices = rng.choice(tmp_len, downsample_len, replace=False)
                train_dataset = Subset(train_dataset, indices=indices)

        train_dataloader = DataLoader(train_dataset, collate_fn=task_to_collator[task_name], batch_size=batch_size)
        valid_dataloader = DataLoader(valid_dataset, collate_fn=task_to_collator[task_name], batch_size=batch_size)
        test_dataloader = DataLoader(test_dataset, collate_fn=task_to_collator[task_name], batch_size=batch_size) \
            if test_dataset is not None else None

        task_to_train_datasets[task_name] = train_dataset
        task_to_train_dataloaders[task_name] = train_dataloader
        task_to_valid_dataloaders[task_name] = valid_dataloader
        task_to_test_dataloaders[task_name] = test_dataloader
        task_to_num_labels[task_name] = num_labels

    print("Dataset lengths: " + " ".join([str(len(dataset)) for dataset in task_to_train_datasets.values()]))

    multitask_train_dataset = MultitaskDataset(task_to_train_datasets)
    multitask_train_sampler = MultitaskBatchSampler(task_to_train_datasets, batch_size)
    multitask_train_collator = MultitaskCollator(task_to_collator)
    multitask_train_dataloader = DataLoader(
        multitask_train_dataset,
        batch_sampler=multitask_train_sampler,
        collate_fn=multitask_train_collator.collator_fn,
    )

    # get the metric function
    task_to_metrics = {}
    for task_name in task_names:
        task_to_metrics[task_name] = load_metric(task_to_metric_name[task_name], task_name)

    return multitask_train_dataloader, task_to_train_dataloaders, task_to_valid_dataloaders, task_to_test_dataloaders, \
        task_to_num_labels, task_to_metrics, config, tokenizer 

def main(args):
    setup_logging_logic()
    logger = get_logger("main")

    # load multitask data
    if args.use_ws_dataset:
        multitask_train_dataloader, task_to_train_loaders, task_to_valid_loaders, task_to_test_loaders, \
        task_to_num_labels, task_to_metrics, transformer_config, transformer_tokenizer = \
            load_ws_task_data(args.ws_task_name, args.model_name_or_path, 
            args.pad_to_max_length, args.max_length, args.batch_size, args.lf_idxes,
            downsample_frac = args.downsample_frac)
    else:
        multitask_train_dataloader, task_to_train_loaders, task_to_valid_loaders, task_to_test_loaders, \
        task_to_num_labels, task_to_metrics, transformer_config, transformer_tokenizer = \
            load_task_data(args.task_names, args.model_name_or_path, 
            args.pad_to_max_length, args.max_length, args.batch_size,
            downsample_frac = args.downsample_frac)
    
    task_names = list(task_to_train_loaders.keys())

    # initialize model
    multitask_model = MultitaskBertForClassification.from_pretrained(
        args.model_name_or_path,
        from_tf=bool(".ckpt" in args.model_name_or_path),
        config = transformer_config,
        tasks = list(task_to_num_labels.keys()),
        num_labels_list = list(task_to_num_labels.values()),
        use_one_predhead = args.use_one_predhead
    )

    device = torch.device(f"cuda:{args.device}" if args.device != "cpu" else "cpu")
    multitask_model.to(device)
    valid_metrics = {}; test_metrics = {}
    for run in range(args.runs):
        # Optimizer
        # Split weights in two groups, one with weight decay and the other not.
        no_decay = ["bias", "LayerNorm.weight"]
        optimizer_grouped_parameters = [
            {
                "params": [p for n, p in multitask_model.named_parameters() if not any(nd in n for nd in no_decay)],
                "weight_decay": args.weight_decay,
            },
            {
                "params": [p for n, p in multitask_model.named_parameters() if any(nd in n for nd in no_decay)],
                "weight_decay": 0.0,
            },
        ]
        optimizer = AdamW(optimizer_grouped_parameters, lr=args.lr)

        # Scheduler and math around the number of training steps.
        trainer_config = { 
            "num_train_epochs": args.epochs,
            "gradient_accumulation_steps": args.gradient_accumulation_steps,
            "max_train_steps": args.max_train_steps,
            "num_warmup_steps": args.num_warmup_steps,

            "save_dir": args.save_dir,
            "save_period": args.save_period,
            "verbosity": args.verbosity,
            
            "monitor": " ".join([args.monitor_mode, args.monitor_metric]),
            "early_stop": args.early_stop
        }
        num_update_steps_per_epoch = math.ceil(len(multitask_train_dataloader) / trainer_config["gradient_accumulation_steps"])
        if trainer_config["max_train_steps"] == -1:
            trainer_config["max_train_steps"] = trainer_config["num_train_epochs"] * num_update_steps_per_epoch
        else:
            trainer_config["num_train_epochs"] = math.ceil(trainer_config["max_train_steps"] / num_update_steps_per_epoch)
        
        lr_scheduler = get_scheduler(
            name=args.lr_scheduler_type,
            optimizer=optimizer,
            num_warmup_steps=trainer_config["num_warmup_steps"]*num_update_steps_per_epoch,
            num_training_steps=trainer_config["max_train_steps"],
        )

        checkpoint_dir = os.path.join("saved", args.save_name + "_" + "_".join(task_names)) # + f"_{run}"
        trainer = MultitaskTrainer(multitask_model, optimizer, lr_scheduler, trainer_config, device, task_to_metrics, 
                                multitask_train_data_loader=multitask_train_dataloader,
                                train_data_loaders = task_to_train_loaders,
                                valid_data_loaders = task_to_valid_loaders,
                                test_data_loaders = task_to_test_loaders if args.use_ws_dataset else None, # stop using test datasets (label not available)
                                task_to_num_labels = task_to_num_labels,
                                checkpoint_dir = checkpoint_dir)
        if not (len(list(os.listdir(checkpoint_dir)))!=0 and args.skip_pretrain):
            trainer.train()
        if args.train_finetune:
            valid_log = {}; test_log = {}
            finetune_tasks = task_names if not args.finetune_target else args.target_tasks
            for task in finetune_tasks:
                finetune_checkpoint_dir = checkpoint_dir + "_finetune_{}".format(task)
                trainer.load_best_model(strict=False)

                # reinitiate optimizer & lr_scheduler
                no_decay = ["bias", "LayerNorm.weight"]
                optimizer_grouped_parameters = [
                    {
                        "params": [p for n, p in multitask_model.named_parameters() if not any(nd in n for nd in no_decay)],
                        "weight_decay": args.weight_decay,
                    },
                    {
                        "params": [p for n, p in multitask_model.named_parameters() if any(nd in n for nd in no_decay)],
                        "weight_decay": 0.0,
                    },
                ]
                optimizer = AdamW(optimizer_grouped_parameters, lr=args.lr)
                # change trainer config to finetune config
                trainer_config['num_train_epochs'] = args.finetune_epochs
                trainer_config["max_train_steps"] = args.max_train_steps
                num_update_steps_per_epoch = math.ceil(len(task_to_train_loaders[task]) / trainer_config["gradient_accumulation_steps"])
                if trainer_config["max_train_steps"] == -1:
                    trainer_config["max_train_steps"] = trainer_config["num_train_epochs"] * num_update_steps_per_epoch
                else:
                    trainer_config["num_train_epochs"] = math.ceil(trainer_config["max_train_steps"] / num_update_steps_per_epoch)
                lr_scheduler = get_scheduler(
                    name=args.lr_scheduler_type,
                    optimizer=optimizer,
                    num_warmup_steps=trainer_config["num_warmup_steps"]*num_update_steps_per_epoch,
                    num_training_steps=trainer_config["max_train_steps"],
                )

                # create fine-tune trainer
                finetune_trainer = MultitaskTrainer(multitask_model, optimizer, lr_scheduler, trainer_config, device, 
                                    task_to_metrics = {task: task_to_metrics[task]}, 
                                    multitask_train_data_loader=multitask_train_dataloader,
                                    train_data_loaders = {task: task_to_train_loaders[task]},
                                    valid_data_loaders = {task: task_to_valid_loaders[task]},
                                    test_data_loaders = task_to_test_loaders if args.use_ws_dataset else None, # stop using test datasets (label not available)
                                    task_to_num_labels = {task: task_to_num_labels[task]},
                                    checkpoint_dir = finetune_checkpoint_dir)
                finetune_trainer.train()

                task_valid_log = finetune_trainer.test(use_valid = True)
                task_test_log = finetune_trainer.test(use_valid = False)
                valid_log.update(task_valid_log)
                test_log.update(task_test_log)
            if args.finetune_target:
                mtl_valid_log = trainer.test(use_valid=True)
                mtl_test_log = trainer.test(use_valid=False) # update log of other tasks
                for key, val in mtl_valid_log.items():
                    if (key == 'loss') or (key not in valid_log):
                        valid_log[key] = val
                for key, val in mtl_test_log.items():
                    if (key == 'loss') or (key not in test_log):
                        test_log[key] = val
        else:
            valid_log = trainer.test(use_valid = True)
            test_log = trainer.test(use_valid = False)
        
        for key, val in valid_log.items():
            if key in valid_metrics:
                valid_metrics[key].append(val)
            else:
                valid_metrics[key] = [val, ]
        
        for key, val in test_log.items():
            if key in test_metrics:
                test_metrics[key].append(val)
            else:
                test_metrics[key] = [val, ]

        # re-initialize model
        multitask_model.load_state_dict(
            torch.load(os.path.join(checkpoint_dir, "model_epoch_0.pth"))['state_dict']
        )
    
    # print training results
    for key, vals in test_metrics.items():
        logger.info("{}: {:.4f} +/- {:.4f}".format(key, np.mean(vals), np.std(vals)))

    # save results into .csv
    file_dir = os.path.join("./results/", args.save_name)
    if not os.path.exists(file_dir):
        os.mkdir(file_dir)
    
    for task_name in task_names:
        # save validation results
        result_datapoint = {
            "Task": task_name, 
            "Trained on": task_names,
        }
        for key, vals in valid_metrics.items():
            if task_name in key:
                result_datapoint[key] = np.mean(vals)
                result_datapoint[key+"_std"] = np.std(vals)
        file_name = os.path.join(file_dir, "{}_{}_valid.csv".format(args.save_name, task_name))
        add_result_to_csv(result_datapoint, file_name)
        # save test results
        result_datapoint = {
            "Task": task_name, 
            "Trained on": task_names,
        }
        for key, vals in test_metrics.items():
            if task_name in key:
                result_datapoint[key] = np.mean(vals)
                result_datapoint[key+"_std"] = np.std(vals)
        file_name = os.path.join(file_dir, "{}_{}_test.csv".format(args.save_name, task_name))
        add_result_to_csv(result_datapoint, file_name)

def add_trainer_arguments(parser):
    parser.add_argument("--epochs", type=int, default=5)
    parser.add_argument("--gradient_accumulation_steps", type=int, default=1)
    parser.add_argument("--max_train_steps", type=int, default=-1)
    parser.add_argument("--num_warmup_steps", type=int, default=0)

    parser.add_argument("--save_dir", type=str, default="./saved/test/")
    parser.add_argument("--save_period", type=int, default=1)
    parser.add_argument("--verbosity", type=int, default=2)

    parser.add_argument("--monitor_mode", type=str, default='off', choices=['min', 'max', 'off'])
    parser.add_argument("--monitor_metric", type=str, default='val_loss')
    parser.add_argument("--early_stop", type=int, default=10)
    return parser

def add_dataloader_arguments(parser):
    parser.add_argument('--pad_to_max_length', default=True)
    parser.add_argument('--max_length', type=int, default=512)
    parser.add_argument('--batch_size', type=int, default=8)
    return parser

def add_finetune_arguments(parser):
    parser.add_argument("--train_finetune", action="store_true")
    parser.add_argument("--skip_pretrain", action="store_true")
    parser.add_argument("--finetune_epochs", type=int, default=5)

    parser.add_argument("--finetune_target", action="store_true")
    parser.add_argument("--target_tasks", nargs='+', default=['mrpc'])
    return parser

if __name__ == "__main__":
    parser = argparse.ArgumentParser()

    parser.add_argument("--task_names", nargs='+', default=['mrpc'])
    parser.add_argument("--model_name_or_path", type=str, default="prajjwal1/bert-mini")
    parser.add_argument("--lr_scheduler_type", type=SchedulerType, default="linear",
        choices=["linear", "cosine", "cosine_with_restarts", "polynomial", "constant", "constant_with_warmup"],
    )

    parser.add_argument("--runs", type=int, default=5)
    parser.add_argument("--device", type=int, default=0)
    parser.add_argument("--lr", type=float, default=5e-5)
    parser.add_argument("--weight_decay", type=float, default=0)

    parser.add_argument("--save_name", type=str, default="test")
    parser.add_argument("--downsample_frac", type=float, default=1.0)

    parser.add_argument("--use_ws_dataset", action="store_true")
    parser.add_argument("--use_one_predhead", action="store_true")
    parser.add_argument("--ws_task_name", type=str, default="youtube")
    parser.add_argument("--lf_idxes", nargs='+', type=int)

    parser = add_trainer_arguments(parser)
    parser = add_dataloader_arguments(parser)
    parser = add_finetune_arguments(parser)
    args = parser.parse_args()

    main(args)