run_segmentation.py

import os
import pdb
import math
import random
import logging
import argparse
from os.path import join as pjoin

import torch
import torch.nn as nn
import torch.distributed as dist
import torch.multiprocessing as mp
import numpy as np
from tqdm import tqdm
from datetime import datetime
from torch.utils.data import DataLoader
from transformers import AutoTokenizer, AdamW, get_linear_schedule_with_warmup

from data.wiki_loader import CrossSegWiki727KDataset, CrossSegWikiSectionDataset
from models.segmentation import CrossSegmentBert

def init_logger(log_file=None, log_file_level=logging.NOTSET):
    log_format = logging.Formatter("[%(asctime)s %(levelname)s] %(message)s")
    logger = logging.getLogger()
    logger.setLevel(logging.INFO)

    if log_file and log_file != '':
        file_handler = logging.FileHandler(log_file)
        file_handler.setLevel(log_file_level)
        file_handler.setFormatter(log_format)
        logger.addHandler(file_handler)
    return logger



def collate_fn(batch):
    results = {}
    results['input_ids'] = torch.tensor([example[0] for example in batch])
    results['token_type_ids'] = torch.tensor([example[1] for example in batch])
    results['attention_mask'] = torch.tensor([example[2] for example in batch])
    results['targets'] = torch.tensor([example[3] for example in batch]).float()

    if len(batch[0]) == 5:
        results['sent_intervals'] = [example[4] for example in batch]
    return results


def eval(model, eval_loader, compute_sparsity=False):
    model.eval()
    device = next(model.parameters()).device

    tp, fp, fn = 0, 0, 0
    total_loss = 0.

    if compute_sparsity:
        sum_masked, sum_entries = 0, 0
    for step, batch in tqdm(enumerate(eval_loader), total=len(eval_loader)):

        batch = {k: v.to(device) for k, v in batch.items()}

        with torch.no_grad():
            output = model(**batch)

        logits, loss = output['logits'], output['loss']

        targets = batch['targets']
        pred = torch.sigmoid(logits).round()
        tp += torch.logical_and(targets == 1, pred == 1).sum().item()
        fp += torch.logical_and(targets == 0, pred == 1).sum().item()
        fn += torch.logical_and(targets == 1, pred == 0).sum().item()
        total_loss += loss.item()

        if compute_sparsity:
            sum_masked += output['sparsity_ratio'][0]
            sum_entries += output['sparsity_ratio'][1]

    try:
        precision = round(tp / (tp + fp), 4)
    except:
        precision = 0.

    try:
        recall = round(tp / (tp + fn), 4)
    except:
        recall = 0.

    f_score = round(tp / (tp + 0.5 * (fp + fn)), 4)
    avg_loss = round(total_loss / len(eval_loader), 4)

    sparsity = 1 - (sum_masked / sum_entries) if compute_sparsity else None

    return precision, recall, f_score, avg_loss, sparsity


class MovingAverage:
    def __init__(self, window_size):
        self.window_size = window_size
        self.values = []
        self.sum = 0

    def process(self, value):
        self.values.append(value)
        self.sum += value
        if len(self.values) > self.window_size:
            self.sum -= self.values.pop(0)
        return float(self.sum) / len(self.values)


def train(rank=None, args=None):

    model = CrossSegmentBert(args)
    model.train()


    # DDP
    if rank is not None:
        device = args.devices[rank]
        print(f"Rank: {rank}, Device: {device}")

        dist.init_process_group(
            backend='nccl',
            init_method='env://',
            world_size=args.world_size,
            rank=rank,
        )

        torch.cuda.set_device(device)
        args.device = torch.device(f'cuda:{device}')
        model.to(args.device)
        model = nn.parallel.DistributedDataParallel(model, device_ids=[device], find_unused_parameters=True)
    else:
        model = model.to(args.device)

    # Initialize datasets
    if args.dataset == 'wiki_727':
        train_set = CrossSegWiki727KDataset(args, 'train')
        dev_set = CrossSegWiki727KDataset(args, 'dev')
    elif args.dataset == 'wiki_section':
        train_set = CrossSegWikiSectionDataset(args, 'train')
        dev_set = CrossSegWikiSectionDataset(args, 'validation')

    # Data loaders
    if rank is not None:
        train_sampler = torch.utils.data.distributed.DistributedSampler(
            train_set,
            num_replicas=args.world_size,
            shuffle=True,
            rank=rank,
        )
        train_loader = DataLoader(
            train_set,
            batch_size=args.batch_size,
            shuffle=False,
            collate_fn=collate_fn,
            num_workers=0,
            pin_memory=True,
            sampler=train_sampler,
        )
    else:
        train_loader = DataLoader(train_set, batch_size=args.batch_size, shuffle=True, collate_fn=collate_fn, num_workers=args.num_workers)

    dev_loader = DataLoader(dev_set, batch_size=args.eval_batch_size, shuffle=False, collate_fn=collate_fn, num_workers=args.num_workers)

    # Optimization
    optimizer = AdamW(model.parameters(), lr=args.lr)
    steps_per_epoch = math.ceil(len(train_loader) / args.grad_accum_steps)
    num_training_steps = (steps_per_epoch * args.num_epochs) if args.num_training_steps == None else args.num_training_steps
    num_warmup_steps = num_training_steps // 10 if args.num_training_steps == None else args.num_training_steps
    lr_scheduler = get_linear_schedule_with_warmup(optimizer, num_warmup_steps=num_warmup_steps, num_training_steps=num_training_steps)
    total_accum_steps = 0

    # Tracking moving average (loss)
    window_size = args.grad_accum_steps * args.report_steps
    moving_loss = MovingAverage(window_size=100)
    best_f_score = 0.

    # File name and f-score for current checkpoints
    kept_checkpoints = []

    for epoch in range(args.num_epochs):
        epoch_loss = 0.

        with tqdm(desc='Training', total=len(train_loader)) as pbar:
            for step, batch in enumerate(train_loader):
                    
                batch = {k: v.to(args.device) for k, v in batch.items()}

                output = model(**batch)
                loss = output['loss']
                loss.backward()

                loss = loss.detach()

                # # Gather loss tensor from all ranks
                # if rank is not None:
                #     tensor_list = [torch.zeros_like(loss) for _ in range(args.world_size)]
                #     loss = dist.all_gather(tensor_list, loss)
                #     print(loss)

                epoch_loss += loss.item()
                avg_loss = moving_loss.process(loss.item())

                if ((step + 1) % args.grad_accum_steps == 0) or (step + 1 == len(train_loader)):
                    optimizer.step()
                    lr_scheduler.step()
                    model.zero_grad()
                    total_accum_steps += 1

                    if total_accum_steps % args.report_steps == 0:
                        lr = optimizer.param_groups[0]['lr']

                        if rank is None or rank == 0:
                            args.logger.info(f'(Step {total_accum_steps}) LR: {lr}, Moving Loss: {avg_loss}')

                    if (total_accum_steps > args.skip_eval_steps) and (total_accum_steps % args.eval_steps == 0):

                        pbar.set_description('Evaluating on Dev Set')
                        precision, recall, f_score, total_loss, _ = eval(model, dev_loader)
                        dev_results = f'(Dev) F1 Score: {f_score}, Loss: {total_loss}'
                        pbar.set_description(dev_results)

                        if rank is None or rank == 0:
                            args.logger.info(dev_results)


                            encoder_name = str(args.encoder).split('/')[-1]
                            save_name = f'{encoder_name}_{args.dataset}_{total_accum_steps}_{round(f_score * 1000)}.pth'
                            kept_checkpoints.append((save_name, f_score))
                            torch.save(model.state_dict(), pjoin(args.ckpt_dir, save_name))
                            args.logger.info(f"Saved Checkpoint at \"{save_name}\"")

                            # Keep only top K checkpoints by f-score
                            if len(kept_checkpoints) > args.num_checkpoints:
                                kept_checkpoints = sorted(kept_checkpoints, key=lambda x: x[1])
                                to_remove = pjoin(args.ckpt_dir, kept_checkpoints[0][0])
                                os.remove(to_remove)
                                args.logger.info(f"Deleted Checkpoint at \"{to_remove}\"")
                                kept_checkpoints = kept_checkpoints[1:] 

                        if rank is not None:
                            dist.barrier()

                        model.train()

                pbar.update(1)


def test(args):
    model = CrossSegmentBert(args)
    model.eval()
    model = model.to(args.device)

    sort_key = lambda x: int(x.split('.')[0].split('_')[-1])
    checkpoints = sorted([f for f in os.listdir(args.ckpt_dir) if f.split('.')[-1] in ['pt', 'pth']], key=sort_key, reverse=True)[:args.num_test]

    # Initialize datasets
    if args.dataset == 'wiki_727':
        test_set = CrossSegWiki727KDataset(args, 'test')
    elif args.dataset == 'wiki_section':
        test_set = CrossSegWikiSectionDataset(args, 'test')

    test_loader = DataLoader(test_set, batch_size=args.eval_batch_size, shuffle=False, collate_fn=collate_fn, num_workers=args.num_workers)

    sum_f_score, sum_precision, sum_recall = 0, 0, 0

    for ckpt in checkpoints:
        args.logger.info(f"\nResults for Checkpoint Loaded at \"{ckpt}\":")

        ckpt_dir = pjoin(args.ckpt_dir, ckpt)

        state_dict = torch.load(ckpt_dir, map_location=args.device)

        # Change key name for saved DDP models
        if all([key.startswith('module.') for key in state_dict.keys()]):
            old_keys = list(state_dict.keys()).copy()
            for key in old_keys:
                state_dict[key.replace('module.', '')] = state_dict.pop(key)

        model.load_state_dict(state_dict)

        precision, recall, f_score, total_loss, sparsity_ratio = eval(model, test_loader, compute_sparsity=args.compute_sparsity)
        ckpt_results = f'(Test) F1 Score: {f_score}, Precision: {precision}, Recall: {recall}, Loss: {total_loss}, Sparsity: {sparsity_ratio}'
        print(ckpt_results)
        args.logger.info(ckpt_results)
        sum_f_score += f_score
        sum_precision += precision
        sum_recall += recall

    f, p, r = sum_f_score / args.num_test, sum_precision / args.num_test, sum_recall / args.num_test
    test_results = f'(Test) Average across {args.num_test} examples - F1 Score: {f:.4}, Precision: {p:.4}, Recall: {r:.4}'
    print(test_results)
    args.logger.info(test_results)



if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument('-device', help='GPU index', default=0)
    parser.add_argument('-encoder', help='Pretrained encoder for the cross-segment attention model', default='bert-base-uncased', type=str)
    parser.add_argument('-mode', help='Train or test the model', default='train', type=str, choices=['train', 'test', 'analysis'])
    parser.add_argument('-high_granularity', help='Use high granularity for wikipedia dataset segmentation', action='store_true')
    parser.add_argument('-seed', help='Random seed', default=888, type=int)
    parser.add_argument('-from_checkpoint', help='Load pretrained model from checkpoint', default=False)
    parser.add_argument('-compute_sparsity', action='store_true', help='Report the sparsity ratio during inference')
    parser.add_argument('-heads_per_pattern', default=1, type=int, help='Number of attention heads (per layer) for each pattern')
    
    # Paths
    parser.add_argument('-ckpt_dir', help='Directory for storing model checkpoints', type=str, default=pjoin('checkpoints', 'cross_seg'))
    parser.add_argument('-results_dir', help='Directory for storing the results', type=str, default=pjoin('results', 'cross_seg'))

    parser.add_argument('-dataset', help='Name of the dataset', default='wiki_section', type=str, choices=['wiki_727', 'wiki_section'])
    parser.add_argument('-context_len', help='Token length of the left and right context (input_len = 2 x context_len + 2)', default=128, type=int)
    parser.add_argument('-pad_context', help='Pad left and right context with [PAD]', action='store_true', default=True) # Based on implementation described in paper

    # Effective batch size = batch_size * grad_accum_steps (e.g. 8 x 8 = 64)
    parser.add_argument('-batch_size', help='Batch size during training', type=int, default=8)
    parser.add_argument('-grad_accum_steps', help='Number of steps for gradient accumulation (Effective batch size = batch_size x grad_accum_steps)', type=int, default=8)
    parser.add_argument('-num_warmup_steps', help='Number of warm-up steps (for lr scheduler)', type=int, default=None)
    parser.add_argument('-num_training_steps', help='Total number of training steps (for lr scheduler)', type=int, default=None)
    parser.add_argument('-num_epochs', help='Max number of epochs to train', type=int, default=3)
    parser.add_argument('-lr', help='Learning rate', type=float, default=1e-5)

    parser.add_argument('-num_workers', help='Number of workers for the dataloaders', type=int, default=0)
    parser.add_argument('-eval_steps', help='Number of accumulated steps before each dev set evaluation', type=int, default=1000)
    parser.add_argument('-skip_eval_steps', help='Number of accumulated steps before evaluation on dev set', type=int, default=0)
    parser.add_argument('-report_steps', help='Number of accumulated steps before printing the training loss', type=int, default=1)
    parser.add_argument('-eval_batch_size', help='Batch size during evaluation', type=int, default=16)
    parser.add_argument('-num_checkpoints', help='Number of checkpoints to keep during training (sorted by validation score)', type=int, default=5)

    parser.add_argument('-hidden_size', help='Hidden size of the output classifier', type=int, default=128)

    parser.add_argument('-num_test', help='Number of checkpoints to evaluate on the test set (sorted by validation score)', type=int, default=3)


    # PAL + Distillation Experiments

    parser.add_argument('-use_patterns', help='Use enhanced attention on the encoder', action='store_true')
    parser.add_argument('-pretrained_embedding', help='Name of the pretrained encoder for loading embeddings', type=str, default=None)

    parser.add_argument('-use_pals', help='Add Projected Attention Layer to the model', action='store_true')
    parser.add_argument('-use_pal_patterns', help='Use enhanced attention on the PAL heads', action='store_true')
    parser.add_argument('-hidden_size_pals', type=int, default=256)
    parser.add_argument('-num_attention_heads_pals', type=int, default=4)

    # Distributed Training
    parser.add_argument('-ddp', help='Use Distributed Data Parallel training', action='store_true')
    parser.add_argument('-local_rank', type=int, default=0)
    parser.add_argument('-devices', nargs="+", type=int)


    args = parser.parse_args()


    # Set random seeds
    random.seed(args.seed)
    np.random.seed(args.seed)
    torch.manual_seed(args.seed)
    torch.cuda.manual_seed(args.seed)


    args.data_dir = pjoin('data', args.dataset)

    # Create directories for storing results and model checkpoints
    if not os.path.exists(args.results_dir):
        os.makedirs(args.results_dir, exist_ok=True)
    if not os.path.exists(args.ckpt_dir):
        os.makedirs(args.ckpt_dir, exist_ok=True)

    # Initialize file for logging
    date = datetime.now().strftime("%Y%m%d")
    log_file = f"{args.encoder}_{args.dataset}_{args.mode}_{date}.log"
    args.logger = init_logger(pjoin(args.results_dir, log_file))
    args.logger.info(args)

    # Map encoder name to HuggingFace model path
    if args.encoder != 'transformer':
        try:
            AutoTokenizer.from_pretrained(args.encoder)
        except:
            raise NotImplementedError(f"Encoder name \"{args.encoder_type}\" not found!") 


    
    


    # Distributed data parallel training
    if args.ddp and args.mode == 'train':
        n_gpus = torch.cuda.device_count()
        for device in args.devices:
            if device >= n_gpus:
                raise Error(f"Device index \"{device}\" out of range for \"{n_gpus}\" available GPUs!")

        os.environ["MASTER_ADDR"] = "localhost"

        if not "MASTER_PORT" in os.environ.keys():
            os.environ["MASTER_PORT"] = "29500"

        args.world_size = len(args.devices)
        mp.spawn(train, nprocs=args.world_size, args=(args,))

    else:

        # Set device
        if torch.cuda.is_available():
            if isinstance(args.device, int) or args.device.isdigit():
                args.device = torch.device(f'cuda:{args.device}')
            else:
                args.device = torch.device(args.device)
        else:
            args.device = torch.device('cpu')

        if args.mode == 'train':
            train(args=args)
        elif args.mode == 'test':
            test(args)
        elif args.mode == 'analysis':
            analyze(args)