Pruning

FINAL_EVALUATION_MODE = False  # During checking we will set this flag to True.

######################### DO NOT CHANGE THIS CELL ##########################
import copy
import pickle

import numpy as np
from IPython.display import clear_output
from tqdm.auto import tqdm

import torch
from torch import nn
from torch.utils.data import Dataset, DataLoader
from torch.optim import SGD

np.random.seed(0)
torch.manual_seed(0)

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

print(f"Using {device} device")

######################### DO NOT CHANGE THIS CELL ##########################

# Function that loads the training and validation data as np.array
def load_data_from_file(x_train_path, y_train_path, x_valid_path, y_valid_path):
    X_train = np.load(x_train_path)
    y_train = np.load(y_train_path)

    X_valid = np.load(x_valid_path)
    y_valid = np.load(y_valid_path)

    return X_train, y_train, X_valid, y_valid


# Dataset class
class InMemDataset(Dataset):
    def __init__(self, xs, ys, device='cpu'):
        super().__init__()
        self.dataset = []
        for i in tqdm(range(len(xs))):
            self.dataset.append((torch.tensor(xs[i]).to(device).float(), torch.tensor(ys[i]).to(device).float() ))

    def __len__(self):
        return len(self.dataset)

    def __getitem__(self, idx):
        return self.dataset[idx]

######################### DO NOT CHANGE THIS CELL ##########################

# Let us load the data and create the dataloaders
X_train, y_train, X_valid, y_valid = load_data_from_file(
    "train_data/X_train.npy",
    "train_data/y_train.npy",
    "valid_data/X_valid.npy",
    "valid_data/y_valid.npy",
)

batch_size = 128

_train = InMemDataset(X_train, y_train, device)

_valid = InMemDataset(X_valid, y_valid, device)

loaders = {
    "train" : DataLoader(_train, batch_size=batch_size, shuffle=True),
    "valid" : DataLoader(_valid, batch_size=batch_size, shuffle=False),
}

######################### DO NOT CHANGE THIS CELL ##########################

# The overall criterion defined in the task statement
def score(mse_loss, sparsity, mse_weight=1.5, sparsity_weight=1.5):

    if type(mse_loss) == np.ndarray:
        mse_loss[mse_loss > 1000] = 1000
    else:
        if mse_loss > 1000:
            mse_loss = 1000

    score = (1 - mse_loss / 1000) ** mse_weight * sparsity**sparsity_weight
    return score

# Ratio of zeroed weights to all weights
def get_sparsity(model):
    total_params = 0
    zero_params = 0

    for name, param in model.named_parameters():
        if "weight" in name or "bias" in name:
            total_params += param.numel()
            zero_params += torch.sum(param == 0).item()

    sparsity = zero_params / total_params
    return sparsity


# Mean squared error (MSE)
def compute_error(model, data_loader):
    model.eval()

    losses = 0
    num_of_el = 0
    with torch.no_grad():
        for x, y in data_loader:
            outputs = model(x)
            num_of_el += x.shape[0] * y.shape[1]
            losses += model.loss(outputs, y, reduction="sum")

    return losses / num_of_el


def points(score):
    def scale(x, lower=0.085, upper=0.95, max_points=1.5):
        scaled = min(max(x, lower), upper)
        return (scaled - lower) / (upper - lower) * max_points
    return scale(score)

######################### DO NOT CHANGE THIS CELL ##########################

# Let us define the architecture of our network
class MLP(nn.Module):
    def __init__(self, *args):
        super().__init__()
        self.flatten = nn.Flatten()
        self.layers = nn.Sequential(
            nn.Linear(128, 1024),
            nn.Sigmoid(),
            nn.Linear(1024, 10),
        )

    def forward(self, x):
        x = self.flatten(x)
        logits = self.layers(x)
        return logits

    def loss(self, input, target, reduction="mean"):
        mse_loss = nn.MSELoss(reduction=reduction)
        return mse_loss(input, target)

######################### DO NOT CHANGE THIS CELL ##########################

# Initialisation of the network weights
def init_weights(m):
    ''' Initialize the weights in the module m.'''
    if isinstance(m, nn.Linear):
        torch.nn.init.xavier_normal_(m.weight)
        m.bias.data.fill_(0.01)


# Function for saving the model weights to a file - remember that your weights file must be named: model_parameters.pkl
def save_parameters(model, file_name="model_parameters.pkl", to_file=True):

    params_to_save = {}
    for name, param in model.named_parameters():
        params_to_save[name] = param.to("cpu")
    
    if not to_file:
        return params_to_save
    
    with open(f"{file_name}", "wb") as f:
        pickle.dump(params_to_save, f)


# Function for loading the model weights from a file
def load_parameters(model, file_name="model_parameters.pkl", from_file=True, params=None):

    if from_file:
        with open(f"{file_name}", "rb") as f:
            params_to_load = pickle.load(f)
    else:
        params_to_load = params
        
    for name, param in model.named_parameters():
        with torch.no_grad():
            param[...] = params_to_load[name].to(device)

######################### DO NOT CHANGE THIS CELL ##########################

# Function for training the model
def train_model(model: nn.Module,
              data_loaders: dict[str, DataLoader],
              num_epochs: int,
              optimizer: torch.optim.Optimizer,
              verbose: bool =True
              ) -> tuple[torch.Tensor, float]:

    """Function for training the model.

    Args:
        model (torch.nn.Module): The neural network to train.
        data_loaders (dict[str, DataLoader]): A dictionary containing the DataLoaders for the training and validation sets.
        num_epochs (int): The number of epochs to train for.
        optimizer (torch.optim.Optimizer): The optimiser used to train the model.
        verbose (bool, optional): If True, shows the training progress.

    Returns:
        tuple[torch.Tensor, float]: A tuple containing the best set of model parameters found during training and the corresponding value of the loss function on the validation set.
    """
    if FINAL_EVALUATION_MODE:
        verbose = False

    best_epoch = None
    best_params = None
    best_val_loss = np.inf

    for epoch in range(num_epochs):
        model.train()
        _iter = 1
        for inputs, targets in data_loaders['train']:
            optimizer.zero_grad()
            outputs = model(inputs)
            loss = model.loss(outputs, targets)
            loss.backward()
            optimizer.step()

            if verbose:
                if _iter % 10 == 0:
                    print(f"Minibatch {_iter:>6}    |  loss {loss.item():>5.2f}  |")

            _iter +=1

        val_loss = compute_error(model, data_loaders["valid"])

        if val_loss < best_val_loss:
            best_epoch = epoch
            best_val_loss = val_loss
            best_params = [copy.deepcopy(p.detach().cpu()) for p in model.parameters()]

        if verbose:
            clear_output(True)
            m = f"After epoch {epoch:>2} | valid loss: {val_loss:>5.2f}"
            print("{0}\n{1}\n{0}".format("-" * len(m), m))

    if best_params is not None:
        if verbose:
            print(f"\nLoading best params on validation set in epoch {best_epoch} with loss {best_val_loss:.2f}")
        with torch.no_grad():
            for param, best_param in zip(model.parameters(), best_params):
                param[...] = best_param

    return best_params, best_val_loss

######################### DO NOT CHANGE THIS CELL ##########################

initial_model = MLP().to(device)
initial_model.apply(init_weights)

optimizer = SGD(
    initial_model.parameters(),
    lr = 0.01,
    momentum = 0.95,
    weight_decay = 0.001)

best_params, best_val_loss = train_model(initial_model, loaders, num_epochs=100, optimizer=optimizer, verbose=True)

loss = compute_error(initial_model, loaders["valid"])
m = f"| Validation loss: {loss:>5.2f} |"
print("{0}\n{1}\n{0}".format("-" * len(m), m))

def starter_pruning_algorithm(model):
    with torch.no_grad():
        model.layers[0].weight[:, 0:2] = 0
    return model

if not FINAL_EVALUATION_MODE:
    # Let us make a deep copy, so as not to change the weights of the trained model
    model_to_prune = copy.deepcopy(initial_model)

    # Let us prune the weights with the example solution
    model_to_prune = starter_pruning_algorithm(model_to_prune)

    # Saving the model parameters (here we changed the file name; you should save as "model_parameters.pkl")
    save_parameters(model_to_prune, "starter_model_parameters.pkl")

    # Now let us see how to load previously saved weights into a newly created model
    new_model = MLP().to(device)
    loss = compute_error(new_model, loaders["valid"])
    print(f"The new model has loss {loss:.3f}")

    # Loading the model parameters
    load_parameters(new_model, "starter_model_parameters.pkl")
    loss = compute_error(new_model, loaders["valid"])
    print(f"After loading the parameters, the model has loss {loss:.3f}")

    mse = compute_error(new_model, loaders["valid"])
    sparsity = get_sparsity(new_model)

    print(f"Model MSE: {mse:.3f} Sparsity: {sparsity:.3f}")
    model_score = score(mse, sparsity)
    print(f"Your model's score is {model_score:.3f}!")
    print(f"Your solution gets {points(model_score):.3f}/1.5 points!")

def your_pruning_algorithm(model):
    # TODO
    pruned_model = starter_pruning_algorithm(model)
    # Saving the model parameters 
    save_parameters(pruned_model, "model_parameters.pkl")
    return pruned_model

model_to_prune = copy.deepcopy(initial_model)
your_pruning_algorithm(model_to_prune)

def evaluate(X_test, y_test):
    """Validator"""
    test_model = MLP().to(device)
    load_parameters(test_model)

    batch_size = 128

    _test = InMemDataset(X_test, y_test, device)
    test_loader = DataLoader(_test, batch_size=batch_size, shuffle=False)

    mse = compute_error(test_model, test_loader)
    sparsity = get_sparsity(test_model)

    print(f"Model had error: {mse:.3f} and sparsity: {sparsity:.3f}")
    model_score = score(mse, sparsity)

    return model_score

if not FINAL_EVALUATION_MODE:    
    model_score = evaluate(X_valid, y_valid)
    print(f"Your solution gets score {model_score:.3f} on validation set.")
    print(f"Your solution gets {points(model_score):.3f}/1.5 points on validation set!")