evaluate_with_classifier.py

from time import time

import torch
import torch.nn as nn
from torch import optim
from torch.optim.lr_scheduler import StepLR

from arguments import parse_args
from dataset import load_dataset
from meter import RunningMeter, BestMeter
from model import Classifier
from utils import (
    save_meter,
    compute_best_metrics,
    compute_classifier_metrics,
    set_all_seeds,
)

# ------------------------------------------------------------------------------

def evaluate_with_classifier(args):
    """
    Evaluating the performance of CPC with a MLP classifier
    :param args: arguments
    :return: None
    """
    results = {}

    print("Percentage of data used: {}%".format(args.data_percentage))

    # Loading the data
    data_loaders, dataset_sizes = load_dataset(args, classifier=True)

    # Creating the model
    model = Classifier(args).to(args.device)

    # Loading pre-trained weights if available
    if args.saved_model is not None:
        model.load_pretrained_weights(args)

    # Optimizer settings
    optimizer = optim.Adam(model.parameters(), lr=args.classifier_lr)
    scheduler = StepLR(optimizer, step_size=25, gamma=0.8)
    criterion = nn.CrossEntropyLoss()

    # Tracking meters
    running_meter = RunningMeter(args=args)
    best_meter = BestMeter()

    for epoch in range(0, args.num_epochs):
        since = time()
        print("Epoch {}/{}".format(epoch, args.num_epochs - 1))
        print("-" * 10)

        # Training
        model, optimizer, scheduler = train(
            model,
            data_loaders["train"],
            criterion,
            optimizer,
            scheduler,
            args,
            epoch,
            dataset_sizes["train"],
            running_meter,
        )

        # Validation
        evaluate(
            model,
            data_loaders["val"],
            args,
            criterion,
            epoch,
            phase="val",
            dataset_size=dataset_sizes["val"],
            running_meter=running_meter,
        )

        # Evaluating on the test data
        evaluate(
            model,
            data_loaders["test"],
            args,
            criterion,
            epoch,
            phase="test",
            dataset_size=dataset_sizes["test"],
            running_meter=running_meter,
        )

        # Saving the logs
        save_meter(args, running_meter, finetune=True)

        # Printing the time taken
        time_elapsed = time() - since
        print(
            "Epoch {} completed in {:.0f}m {:.0f}s".format(
                epoch, time_elapsed // 60, time_elapsed % 60
            )
        )

    # Computing the best metrics
    best_meter = compute_best_metrics(running_meter, best_meter, classifier=True)
    running_meter.update_best_meter(best_meter)
    save_meter(args, running_meter, finetune=True)

    # Printing the best metrics corresponding to the highest validation
    # F1-score
    best_meter.display()
    # Creating the results dictionary with relevant metrics
    results = {
        "val_loss": best_meter.loss["val"],
        "test_accuracy": best_meter.accuracy["test"],
        "test_f1_score": best_meter.f1_score["test"],
    }

    return results


def train(
    model,
    data_loader,
    criterion,
    optimizer,
    scheduler,
    args,
    epoch,
    dataset_size,
    running_meter,
):
    # Setting the model to training mode
    model.train()

    # Freeze encoder layers
    if args.learning_schedule == "last_layer":
        model.freeze_encoder_layers()

    # To track the loss and other metrics
    running_loss = 0.0
    actual_labels = []
    pred_labels = []

    # Iterating over the data
    for inputs, labels in data_loader:
        inputs = inputs.float().to(args.device)
        labels = labels.long().to(args.device)

        optimizer.zero_grad()

        with torch.set_grad_enabled(True):
            outputs = model(inputs)
            _, preds = torch.max(outputs, 1)
            loss = criterion(outputs, labels)

            loss.backward()
            optimizer.step()

        # Appending predictions and loss
        running_loss += loss.item() * inputs.size(0)
        actual_labels.extend(labels.cpu().data.numpy())
        pred_labels.extend(preds.cpu().data.numpy())

    scheduler.step()

    # Statistics
    loss = running_loss / dataset_size
    _ = compute_classifier_metrics(
        actual_labels, pred_labels, "train", running_meter, loss, epoch
    )

    return model, optimizer, scheduler


def evaluate(
    model, data_loader, args, criterion, epoch, phase, dataset_size, running_meter
):
    # Setting the model to eval mode
    model.eval()

    # To track the loss and other metrics
    running_loss = 0.0
    actual_labels = []
    pred_labels = []

    # Iterating over the data
    for inputs, labels in data_loader:
        inputs = inputs.float().to(args.device)
        labels = labels.long().to(args.device)

        with torch.set_grad_enabled(False):
            outputs = model(inputs)
            _, preds = torch.max(outputs, 1)
            loss = criterion(outputs, labels)

        # Appending predictions and loss
        running_loss += loss.item() * inputs.size(0)
        actual_labels.extend(labels.cpu().data.numpy())
        pred_labels.extend(preds.cpu().data.numpy())

    # Statistics
    loss = running_loss / dataset_size
    _ = compute_classifier_metrics(
        actual_labels, pred_labels, phase, running_meter, loss, epoch
    )

    return


# ------------------------------------------------------------------------------
if __name__ == "__main__":
    args = parse_args()
    set_all_seeds(args)
    print(args)

    evaluate_with_classifier(args=args)

    print("------ Evaluation complete! ------")