nlp_centroid_eval.py

import json
import logging
from pathlib import Path
from typing import Union

import hydra
import numpy as np
import torch
import yaml
from omegaconf import OmegaConf
from torch.utils.data import DataLoader

from src.data import ag_news
from src.data.dataset import DownstreamDataset
from src.eval_utils import centroid_eval
from src.model import CentroidClassifier, ContrastiveFastText, SupervisedFastText


def convert_vectors(
        data_loader: torch.utils.data.DataLoader, model: Union[ContrastiveFastText, SupervisedFastText],
        device: torch.device, normalized: bool
) -> tuple:
    """
    Convert experiment to feature representations.
    :param data_loader: Tata loader for raw experiment.
    :param model: Pre-trained model.
    :param device: PyTorch's device instance.
    :param normalized: Whether normalize the feature representation or not.

    :return: Tuple of tensors: features and labels.
    """

    new_X = []
    new_y = []
    model.eval()

    with torch.no_grad():
        for x_batches, y_batches, offsets in data_loader:
            fs = model(x_batches.to(device), offsets.to(device))
            if normalized:
                fs = torch.nn.functional.normalize(fs, p=2, dim=1)
            new_X.append(fs)
            new_y.append(y_batches)

    X = torch.cat(new_X)
    y = torch.cat(new_y)

    return X, y


@hydra.main(config_path="conf", config_name="nlp_eval_config")
def main(cfg: OmegaConf):
    logger = logging.getLogger(__name__)
    logger.setLevel(logging.INFO)
    stream_handler = logging.StreamHandler()
    stream_handler.setLevel(logging.INFO)
    stream_handler.terminator = ""
    logger.addHandler(stream_handler)

    seed = cfg["experiment"]["seed"]
    torch.manual_seed(seed)
    torch.backends.cudnn.deterministic = True
    torch.backends.cudnn.benchmark = False

    use_cuda = torch.cuda.is_available()

    if use_cuda:
        device_id = cfg["experiment"]["gpu_id"] % torch.cuda.device_count()
        device = torch.device(device_id)
    else:
        device = torch.device("cpu")

    logger_line = "Using {}".format(device)
    logger.info(logger_line)

    # load self-sup training's config
    weights_path = Path(cfg["experiment"]["target_weight_file"])
    weight_name = weights_path.name
    logger.info("Evaluation by using {}".format(weight_name))

    pre_train_config_path = weights_path.parent / ".hydra" / "config.yaml"
    with open(pre_train_config_path) as f:
        pre_train_conf = yaml.load(f, Loader=yaml.FullLoader)

    # initialise data loaders
    training_dataset, validation_dataset = ag_news.get_train_val_datasets(
        root=Path.home() / "pytorch_datasets",
        min_freq=cfg["dataset"]["min_freq"],
    )
    vocab_size = training_dataset.vocab_size
    mask_ratio = pre_train_conf["dataset"]["mask_ratio"]
    aug_type = pre_train_conf["dataset"]["augmentation_type"]

    if aug_type == "erase":
        replace_data = None
    else:
        replace_data = np.load(pre_train_conf["dataset"]["replace_data"])
        assert len(replace_data) == vocab_size

    training_data_loader = DataLoader(
        training_dataset,
        batch_size=cfg["experiment"]["batches"],
        shuffle=False,
        collate_fn=ag_news.CollateSupervised(mask_ratio, replace_data, np.random.RandomState(seed), aug_type),
        drop_last=False
    )

    validation_data_loader = DataLoader(
        validation_dataset,
        batch_size=cfg["experiment"]["batches"],
        shuffle=False,
        collate_fn=ag_news.collate_eval_batch,
        drop_last=False
    )

    num_classes = len(np.unique(validation_dataset.targets))

    logger.info("#train: {}, #val: {}".format(len(training_dataset), len(validation_dataset)))

    assert pre_train_conf["parameter"]["algorithm"] in ["self_supervised", "supervised"]
    is_self_supervised_pretrain = pre_train_conf["parameter"]["algorithm"] == "self_supervised"
    d = pre_train_conf["architecture"]["embedding_dim"]
    if is_self_supervised_pretrain:
        model = ContrastiveFastText(
            num_embeddings=vocab_size,
            embedding_dim=d,
            num_last_hidden_units=d,
            with_projection_head=True
        ).to(device)
    else:
        model = SupervisedFastText(
            num_embeddings=vocab_size,
            embedding_dim=d
        ).to(device)


    state_dict = torch.load(weights_path)
    state_dict = {k.replace("module.", ""): v for k, v in state_dict.items()}

    # load weights trained on self-supervised task
    if use_cuda:
        model.load_state_dict(state_dict, strict=False)
    else:
        model.load_state_dict(state_dict, strict=False, map_location=device)

    # remove projection head or not
    if not cfg["experiment"]["use_projection_head"]:
        if is_self_supervised_pretrain:
            model.g = torch.nn.Identity()
        else:
            model.f.fc = torch.nn.Identity()

    # create feature representations and classifier
    x, y = convert_vectors(training_data_loader, model, device, normalized=True)
    downstream_training_dataset = DownstreamDataset(x, y)

    classifier = CentroidClassifier(
        weights=CentroidClassifier.create_weights(downstream_training_dataset, num_classes=num_classes).to(device)
    )

    # create data_loader for centroids classifier's input
    x, y = convert_vectors(validation_data_loader, model, device, normalized=True)
    downstream_val_dataset = DownstreamDataset(x, y)

    downstream_training_data_loader = DataLoader(
        dataset=downstream_training_dataset, batch_size=cfg["experiment"]["batches"], shuffle=False,
    )
    downstream_val_data_loader = DataLoader(
        dataset=downstream_val_dataset, batch_size=cfg["experiment"]["batches"], shuffle=False,
    )

    top_k = min(cfg["experiment"]["top_k"], num_classes)
    train_acc, train_top_k_acc = centroid_eval(downstream_training_data_loader, device, classifier, top_k)
    val_acc, val_top_k_acc = centroid_eval(downstream_val_data_loader, device, classifier, top_k)

    classification_results = {}
    classification_results[weight_name] = {
        "train_acc": train_acc,
        "train_top_{}_acc".format(top_k): train_top_k_acc,
        "val_acc": val_acc,
        "val_top_{}_acc".format(top_k): val_top_k_acc
    }
    logger.info("train acc: {}, val acc: {}".format(train_acc, val_acc))

    # save evaluation metric
    fname = cfg["experiment"]["classification_results_json_fname"]

    with open(fname, "w") as f:
        json.dump(classification_results, f)


if __name__ == "__main__":
    main()