Self-Supervised Learning (A Classifier in Two Minutes)

FINAL_EVALUATION_MODE = False

import os
import random
import warnings

import gdown
import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F
from sklearn.metrics import accuracy_score
from torch.utils.data import Dataset

warnings.filterwarnings("ignore")

# DO NOT CHANGE THESE VALUES
DEVICE = torch.device("cuda")
MODEL_FP = "./encoder.pt"

# THESE VALUES MAY BE CHANGED
SEED =  42
EPOCHS = 200
N_CLASS = 6
N_CHANNEL = 3
N_LENGTH = 206
PROJECTION_DIM = 64
LOGISTIC_BATCH_SIZE = 64
LOGISTIC_EPOCHS = 10
FINETUNE_SEED = 42

if not FINAL_EVALUATION_MODE:
    ! gdown https://drive.google.com/uc?id=1RxxsAQ6memLe4B0fePJD4cV_SIuRmhmw
    # ! gdown https://drive.google.com/uc?id=1K7873pSrnIEvQe4HxhvD0S6Ur-7bEmfm
    # ! gdown https://drive.google.com/uc?id=16ynUNF1WicH37L5iJszsfvVc-NJPzz8y
    ! unzip self_supervised.zip

class CustomTensorDataset(Dataset):
    """TensorDataset with support of transforms."""
    def __init__(self, data, transform_A=None, transform_B=None):
        assert all(data[0].shape[0] == item.shape[0] for item in data)
        self.data = data
        self.transform_A = transform_A
        self.transform_B = transform_B

    def __getitem__(self, index):
        x = self.data[0][index]

        if self.transform_A:
            x1 = self.transform_A(x)
        else:
            x1 = x
        if self.transform_B:
            x2 = self.transform_B(x)
        else:
            x2 = x
        y = self.data[1][index]

        return torch.tensor(x1).float(), torch.tensor(x2).float(), torch.tensor(y)

    def __len__(self):
        return self.data[0].shape[0]

def setup_seed(seed=42):
    """Setup seed for the reproducipility"""
    torch.manual_seed(seed)
    torch.cuda.manual_seed_all(seed)
    np.random.seed(seed)
    random.seed(seed)
    torch.backends.cudnn.deterministic = True

# Improve the class below!

class SimpleEncoder(nn.Module):
    def __init__(self, projection_dim, n_channel, n_length=240):
        super(SimpleEncoder, self).__init__()
        self.n_features = n_channel * n_length
        self.projector = nn.Sequential(
            nn.Linear(self.n_features, self.n_features // 2),
            nn.ReLU(),
            nn.Linear(self.n_features // 2, projection_dim),
        )

    def forward(self, x_i):
        h_i = x_i.flatten(start_dim=1, end_dim=-1)
        z_i = self.projector(h_i)
        return h_i, z_i

    def save_model(self, model):
        torch.save(model.state_dict(), MODEL_FP)

# Improve the class below!

class MLPClassifier(nn.Module):
    def __init__(self, n_features, n_classes):
        super(MLPClassifier, self).__init__()
        n_dim = n_features // n_classes // 2 * n_classes

        self.model = nn.Sequential(
            nn.Linear(n_features, n_dim),
            nn.ReLU(),
            nn.Linear(n_dim, n_classes)
        )

    def forward(self, x):
        return self.model(x)

# Improve the function below!

def train(unlabelled_train_loader, encoder_model):
    """Train the model (encoder) on unlabelled data."""
    optimizer = None  # TODO
    criterion = None  # TODO
    loss_epoch = []
    for epoch in range(LOGISTIC_EPOCHS):
        for step, (x_i, x_j, _) in enumerate(unlabelled_train_loader):
            optimizer.zero_grad()
            x_i = x_i.to(DEVICE)

            # TODO

            loss = criterion(...)
            loss.backward()
            optimizer.step()
            loss_epoch.append(loss.item())
    mean_loss = sum(loss_epoch) / len(loss_epoch)
    return mean_loss

if not FINAL_EVALUATION_MODE:
    X_train_big = torch.load("train_x_big.pt")
    encoder_model = SimpleEncoder(
        projection_dim=PROJECTION_DIM,
        n_channel=N_CHANNEL,
        n_length=N_LENGTH
    )

    encoder_model = encoder_model.to(DEVICE)
    encoder_model.save_model(encoder_model)
    # TODO: train(train_loader, model)

# Improve the function below!

def finetune(labelled_data_loader, encoder, classifier):
    """Fine-tune the models (encoder, classifier) on labelled data."""
    optimizer = torch.optim.Adam(classifier.parameters(), lr=3e-4)
    criterion = torch.nn.CrossEntropyLoss()
    for epoch in range(LOGISTIC_EPOCHS):
        loss_epoch = []
        encoder.train()
        classifier.train()
        for n_batch, (x, x_aug, y) in enumerate(labelled_data_loader):
            optimizer.zero_grad()
            x = x.to(DEVICE)
            y = y.to(DEVICE)
            with torch.no_grad():
                h, z = encoder(x)
            output = classifier(h).to(DEVICE)
            y = y.squeeze(-1).type(torch.LongTensor).to(DEVICE)
            loss = criterion(output, y)
            loss.backward()
            optimizer.step()
            loss_epoch.append(loss.item())
        mean_loss = sum(loss_epoch) / len(loss_epoch)

    return mean_loss

# Improve the function below!

def finetune_and_predict(X_train_small, y_train_small, X_test, model_path):
    """During testing we will run this function with the test data.
    Do not change the signature of this function, i.e. the number, names and order of its arguments.

    Arguments:
    X_train_small -- tensor of training data of shape (N, 1, 3, 206), where N is the number of independent measurements in the set,
    y_train_small -- vector of labels for the training set, of length N,
    X_test -- tensor of test data of shape (M, 1, 3, 206), where M is the number of independent measurements in the set,

    Return:
    The function should return a vector of length M of predicted labels for the set X_test.
    """
    setup_seed(FINETUNE_SEED)

    train_dataset = CustomTensorDataset((X_train_small, y_train_small))

    train_loader = torch.utils.data.DataLoader(
        train_dataset,
        batch_size=LOGISTIC_BATCH_SIZE,
        shuffle=True,
        drop_last=True,
    )

    encoder_model = SimpleEncoder(
        projection_dim=PROJECTION_DIM,
        n_channel=N_CHANNEL,
        n_length=N_LENGTH
    )

    encoder_model.load_state_dict(torch.load(model_path, map_location=DEVICE.type))
    encoder_model = encoder_model.to(DEVICE)
    n_classes = N_CLASS
    classifier = MLPClassifier(encoder_model.n_features, n_classes)
    classifier = classifier.to(DEVICE)
    mean_loss_train = finetune(train_loader, encoder_model, classifier)
    x = torch.tensor(X_test).to(DEVICE)
    h, z = encoder_model(x)
    output = classifier(h).to(DEVICE)
    predicted = output.argmax(1)
    assert len(predicted.shape) == 1
    assert len(predicted) == len(X_test)
    assert predicted.dtype == torch.int64
    return predicted

X_val = torch.load("val_x.pt")
y_val = torch.load("val_y.pt")
if not FINAL_EVALUATION_MODE:
    X_train_small = torch.load("train_x_small.pt")
    y_train_small = torch.load("train_y_small.pt")
    pred = finetune_and_predict(X_train_small, y_train_small, X_val, MODEL_FP)
    print("Accuracy:", accuracy_score(y_val, pred.cpu().numpy()))