train/finetune/run_classifier_siamese.py

"""
This script provides an example to wrap TencentPretrain for classification with siamese network.
"""
import sys
import os
import random
import argparse
import collections
import torch
import torch.nn as nn

tencentpretrain_dir = os.path.abspath(os.path.join(os.path.dirname(__file__), ".."))
sys.path.append(tencentpretrain_dir)

from tencentpretrain.embeddings import *
from tencentpretrain.encoders import *
from tencentpretrain.targets import *
from tencentpretrain.utils.vocab import Vocab
from tencentpretrain.utils.constants import *
from tencentpretrain.utils import *
from tencentpretrain.utils.optimizers import *
from tencentpretrain.utils.config import load_hyperparam
from tencentpretrain.utils.seed import set_seed
from tencentpretrain.utils.logging import init_logger
from tencentpretrain.utils.misc import pooling
from tencentpretrain.model_saver import save_model
from tencentpretrain.opts import finetune_opts, tokenizer_opts
from finetune.run_classifier import count_labels_num, build_optimizer


class SiameseClassifier(nn.Module):
    def __init__(self, args):
        super(SiameseClassifier, self).__init__()
        self.embedding = Embedding(args)
        for embedding_name in args.embedding:
            tmp_emb = str2embedding[embedding_name](args, len(args.tokenizer.vocab))
            self.embedding.update(tmp_emb, embedding_name)
        self.encoder = DualEncoder(args)

        self.classifier = nn.Linear(4 * args.stream_0["hidden_size"], args.labels_num)
        self.pooling_type = args.pooling

    def forward(self, src, tgt, seg):
        """
        Args:
            src: [batch_size x seq_length]
            tgt: [batch_size]
            seg: [batch_size x seq_length]
        """
        # Embedding.
        emb = self.embedding(src, seg)
        # Encoder.
        output = self.encoder(emb, seg)
        # Target.
        features_0, features_1 = output
        features_0 = pooling(features_0, seg[0], self.pooling_type)
        features_1 = pooling(features_1, seg[1], self.pooling_type)

        vectors_concat = []

        # concatenation
        vectors_concat.append(features_0)
        vectors_concat.append(features_1)
        # difference:
        vectors_concat.append(torch.abs(features_0 - features_1))
        # multiplication:
        vectors_concat.append(features_0 * features_1)

        features = torch.cat(vectors_concat, 1)

        logits = self.classifier(features)

        if tgt is not None:
            loss = nn.NLLLoss()(nn.LogSoftmax(dim=-1)(logits), tgt.view(-1))
            return loss, logits
        else:
            return None, logits


def load_or_initialize_parameters(args, model):
    if args.pretrained_model_path is not None:
        # Initialize with pretrained model.
        state_dict = torch.load(args.pretrained_model_path, map_location="cpu")
        load_siamese_weights = False
        for key in state_dict.keys():
            if key.find("embedding_0") != -1:
                load_siamese_weights = True
                break
        if not load_siamese_weights:
            siamese_state_dict = collections.OrderedDict()
            for key in state_dict.keys():
                if key.split('.')[0] == "embedding":
                    siamese_state_dict["embedding.embedding_0." + ".".join(key.split('.')[1:])] = state_dict[key]
                    siamese_state_dict["embedding.embedding_1." + ".".join(key.split('.')[1:])] = state_dict[key]
                if key.split('.')[0] == "encoder":
                    siamese_state_dict["encoder.encoder_0." + ".".join(key.split('.')[1:])] = state_dict[key]
                    siamese_state_dict["encoder.encoder_1." + ".".join(key.split('.')[1:])] = state_dict[key]
            model.load_state_dict(siamese_state_dict, strict=False)
        else:
            model.load_state_dict(state_dict, strict=False)
    else:
        # Initialize with normal distribution.
        for n, p in list(model.named_parameters()):
            if "gamma" not in n and "beta" not in n:
                p.data.normal_(0, 0.02)


def batch_loader(batch_size, src, tgt, seg):
    instances_num = tgt.size()[0]
    src_a, src_b = src
    seg_a, seg_b = seg
    for i in range(instances_num // batch_size):
        src_a_batch = src_a[i * batch_size : (i + 1) * batch_size, :]
        src_b_batch = src_b[i * batch_size : (i + 1) * batch_size, :]
        tgt_batch = tgt[i * batch_size : (i + 1) * batch_size]
        seg_a_batch = seg_a[i * batch_size : (i + 1) * batch_size, :]
        seg_b_batch = seg_b[i * batch_size : (i + 1) * batch_size, :]
        yield (src_a_batch, src_b_batch), tgt_batch, (seg_a_batch, seg_b_batch)
    if instances_num > instances_num // batch_size * batch_size:
        src_a_batch = src_a[instances_num // batch_size * batch_size :, :]
        src_b_batch = src_b[instances_num // batch_size * batch_size :, :]
        tgt_batch = tgt[instances_num // batch_size * batch_size :]
        seg_a_batch = seg_a[instances_num // batch_size * batch_size :, :]
        seg_b_batch = seg_b[instances_num // batch_size * batch_size :, :]
        yield (src_a_batch, src_b_batch), tgt_batch, (seg_a_batch, seg_b_batch)


def read_dataset(args, path):
    dataset, columns = [], {}
    with open(path, mode="r", encoding="utf-8") as f:
        for line_id, line in enumerate(f):
            if line_id == 0:
                for i, column_name in enumerate(line.rstrip("\r\n").split("\t")):
                    columns[column_name] = i
                continue
            line = line.rstrip("\r\n").split("\t")
            tgt = int(line[columns["label"]])

            text_a, text_b = line[columns["text_a"]], line[columns["text_b"]]
            src_a = args.tokenizer.convert_tokens_to_ids([CLS_TOKEN] + args.tokenizer.tokenize(text_a) + [SEP_TOKEN])
            src_b = args.tokenizer.convert_tokens_to_ids([CLS_TOKEN] + args.tokenizer.tokenize(text_b) + [SEP_TOKEN])
            seg_a = [1] * len(src_a)
            seg_b = [1] * len(src_b)
            PAD_ID = args.tokenizer.convert_tokens_to_ids([PAD_TOKEN])[0]

            if len(src_a) >= args.seq_length:
                src_a = src_a[:args.seq_length]
                seg_a = seg_a[:args.seq_length]
            while len(src_a) < args.seq_length:
                src_a.append(PAD_ID)
                seg_a.append(0)

            if len(src_b) >= args.seq_length:
                src_b = src_b[:args.seq_length]
                seg_b = seg_b[:args.seq_length]
            while len(src_b) < args.seq_length:
                src_b.append(PAD_ID)
                seg_b.append(0)

            dataset.append(((src_a, src_b), tgt, (seg_a, seg_b)))

    return dataset


def train_model(args, model, optimizer, scheduler, src_batch, tgt_batch, seg_batch):
    model.zero_grad()

    src_a_batch, src_b_batch = src_batch
    seg_a_batch, seg_b_batch = seg_batch

    src_a_batch = src_a_batch.to(args.device)
    src_b_batch = src_b_batch.to(args.device)

    tgt_batch = tgt_batch.to(args.device)

    seg_a_batch = seg_a_batch.to(args.device)
    seg_b_batch = seg_b_batch.to(args.device)

    loss, _ = model((src_a_batch, src_b_batch), tgt_batch, (seg_a_batch, seg_b_batch))

    if torch.cuda.device_count() > 1:
        loss = torch.mean(loss)

    if args.fp16:
        with args.amp.scale_loss(loss, optimizer) as scaled_loss:
            scaled_loss.backward()
    else:
        loss.backward()

    optimizer.step()
    scheduler.step()

    return loss


def evaluate(args, dataset):
    src_a = torch.LongTensor([example[0][0] for example in dataset])
    src_b = torch.LongTensor([example[0][1] for example in dataset])
    tgt = torch.LongTensor([example[1] for example in dataset])
    seg_a = torch.LongTensor([example[2][0] for example in dataset])
    seg_b = torch.LongTensor([example[2][1] for example in dataset])

    batch_size = args.batch_size

    correct = 0
    # Confusion matrix.
    confusion = torch.zeros(args.labels_num, args.labels_num, dtype=torch.long)

    args.model.eval()

    for i, (src_batch, tgt_batch, seg_batch) in enumerate(batch_loader(batch_size, (src_a, src_b), tgt, (seg_a, seg_b))):

        src_a_batch, src_b_batch = src_batch
        seg_a_batch, seg_b_batch = seg_batch

        src_a_batch = src_a_batch.to(args.device)
        src_b_batch = src_b_batch.to(args.device)

        tgt_batch = tgt_batch.to(args.device)

        seg_a_batch = seg_a_batch.to(args.device)
        seg_b_batch = seg_b_batch.to(args.device)

        with torch.no_grad():
            _, logits = args.model((src_a_batch, src_b_batch), None, (seg_a_batch, seg_b_batch))
        pred = torch.argmax(nn.Softmax(dim=1)(logits), dim=1)
        gold = tgt_batch
        for j in range(pred.size()[0]):
            confusion[pred[j], gold[j]] += 1
        correct += torch.sum(pred == gold).item()

    args.logger.debug("Confusion matrix:")
    args.logger.debug(confusion)
    args.logger.debug("Report precision, recall, and f1:")

    eps = 1e-9
    for i in range(confusion.size()[0]):
        p = confusion[i, i].item() / (confusion[i, :].sum().item() + eps)
        r = confusion[i, i].item() / (confusion[:, i].sum().item() + eps)
        f1 = 2 * p * r / (p + r + eps)
        args.logger.debug("Label {}: {:.3f}, {:.3f}, {:.3f}".format(i, p, r, f1))

    args.logger.info("Acc. (Correct/Total): {:.4f} ({}/{}) ".format(correct / len(dataset), correct, len(dataset)))
    return correct / len(dataset), confusion


def main():
    parser = argparse.ArgumentParser(formatter_class=argparse.ArgumentDefaultsHelpFormatter)

    finetune_opts(parser)

    tokenizer_opts(parser)

    args = parser.parse_args()

    # Load the hyperparameters from the config file.
    args = load_hyperparam(args)

    set_seed(args.seed)

    # Count the number of labels.
    args.labels_num = count_labels_num(args.train_path)

    # Build tokenizer.
    args.tokenizer = str2tokenizer[args.tokenizer](args)

    # Build classification model.
    model = SiameseClassifier(args)

    # Load or initialize parameters.
    load_or_initialize_parameters(args, model)

    # Get logger.
    args.logger = init_logger(args)

    args.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    model = model.to(args.device)

    # Training phase.
    trainset = read_dataset(args, args.train_path)
    instances_num = len(trainset)
    batch_size = args.batch_size

    args.train_steps = int(instances_num * args.epochs_num / batch_size) + 1

    args.logger.info("Batch size: {}".format(batch_size))
    args.logger.info("The number of training instances: {}".format(instances_num))

    optimizer, scheduler = build_optimizer(args, model)

    if args.fp16:
        try:
            from apex import amp
        except ImportError:
            raise ImportError("Please install apex from https://www.github.com/nvidia/apex to use fp16 training.")
        model, optimizer = amp.initialize(model, optimizer, opt_level=args.fp16_opt_level)
        args.amp = amp

    if torch.cuda.device_count() > 1:
        args.logger.info("{} GPUs are available. Let's use them.".format(torch.cuda.device_count()))
        model = torch.nn.DataParallel(model)
    args.model = model

    total_loss, result, best_result = 0.0, 0.0, 0.0

    args.logger.info("Start training.")

    for epoch in range(1, args.epochs_num + 1):
        random.shuffle(trainset)
        src_a = torch.LongTensor([example[0][0] for example in trainset])
        src_b = torch.LongTensor([example[0][1] for example in trainset])
        tgt = torch.LongTensor([example[1] for example in trainset])
        seg_a = torch.LongTensor([example[2][0] for example in trainset])
        seg_b = torch.LongTensor([example[2][1] for example in trainset])

        model.train()
        for i, (src_batch, tgt_batch, seg_batch) in enumerate(batch_loader(batch_size, (src_a, src_b), tgt, (seg_a, seg_b))):
            loss = train_model(args, model, optimizer, scheduler, src_batch, tgt_batch, seg_batch)
            total_loss += loss.item()
            if (i + 1) % args.report_steps == 0:
                args.logger.info("Epoch id: {}, Training steps: {}, Avg loss: {:.3f}".format(epoch, i + 1, total_loss / args.report_steps))
                total_loss = 0.0

        result = evaluate(args, read_dataset(args, args.dev_path))
        if result[0] > best_result:
            best_result = result[0]
            save_model(model, args.output_model_path)

    # Evaluation phase.
    if args.test_path is not None:
        args.logger.info("Test set evaluation.")
        if torch.cuda.device_count() > 1:
            args.model.module.load_state_dict(torch.load(args.output_model_path))
        else:
            args.model.load_state_dict(torch.load(args.output_model_path))
        evaluate(args, read_dataset(args, args.test_path))


if __name__ == "__main__":
    main()