<jupyter_start><jupyter_text>Homework 5Starting from the implementation contained within the notebook `05-pruning.ipynb`, extend the  `magnitude_pruning` function to allow for incremental (iterative) pruning. Right now, the magnitude_pruning routine is thought for one-shot pruning. If you try pruning one more time, you'll notice that it will not work as there's no way to communicate to the future calls of magnitude_pruning to ignore the parameters which have already been pruned. Find a way to enhance the routine s.t. it can effectively prune networks in a sequential fashion (i.e., if we passed an MLP already pruned of 20% of its parameters, we want to prune *another* 20% of parameters).Hint: use the mask!<jupyter_code>import torch

sys.path.append("../labs")
from scripts import mnist, train_utils, architectures, train
from scripts.train_utils import accuracy, AverageMeter
layers = [
    {"n_in": 784, "n_out": 16, "batchnorm": False},
    {"n_out": 32, "batchnorm": True},
    {"n_out": 64, "batchnorm": True},
    {"n_out": 10, "batchnorm": True}
]
net = architectures.MLPCustom(layers)
print(net)
def magnitude_pruning(model, pruning_rate, layers_to_prune=["1", "4", "7", "10"]):
    # 1. vectorize distribution of abs(parameter)
    '''
    model.named_parameters returns a tuple:
      element 0 is the name of the parameter
      element 1 is the parameter itself
    any(<list of booleans>)
    checks whether any of the conditions inside the list is True.
    [l in pars[0] for l in layers_to_prune] checks whether "1", "4", "7", and "10" are inside the parameter name
    hence, there will be up to 1 True in the corresponding list: in that case, any(...) returns True
    '''
    params_to_prune = [pars[1] for pars in model.named_parameters() if any([l in pars[0] for l in layers_to_prune])]
    flat = torch.cat([pars.abs().flatten() for pars in params_to_prune], dim=0)

    # 2. sort this distribution
    flat = flat.sort()[0]

    # 3. obtain the threshold
    position = int(pruning_rate * flat.shape[0])
    thresh = flat[position]

    # 4. binarize the parameters & 5. compose these booleans into the mask
    # 6. obtain the new structure of parameters
    '''
    I do this process with a for cycle instead of a list comprehension for clarity
    * if the layer is a layer to prune → populate the mask with 1s and 0s
    * otherwise → just populate the mask with ones
    By doing so, I can immediately apply the mask to the model as well...
    '''
    mask = []
    for pars in model.named_parameters():
        if any([l in pars[0] for l in layers_to_prune]):
            m = torch.where(pars[1].abs() >= thresh, 1, 0)
            mask.append(m)
            pars[1].data *= m
        else:
            mask.append(torch.ones_like(pars[1]))
            # no need to multiply as it's all 1s

    # 7. what do we need to return?
    return mask
mask = magnitude_pruning(net, .2)
def number_of_ones_in_mask(mask):
    return sum([m.sum().item() for m in mask]) / sum([m.numel() for m in mask])
# Proportion of parameters which have survived the pruning
print("Number of ones in mask:", number_of_ones_in_mask(mask), "\n")
# I'm pruning the first 20 parameters of the network and the mask to see if it actually works
print(mask[0][0,:20], "\n", next(net.parameters())[0,:20])
trainloader, testloader, _, _ = mnist.get_data(data_root="../labs/datasets/")
loss_fn = torch.nn.CrossEntropyLoss()
# load pretrained model
state_dict = torch.load("../labs/models_push/mlp_custom_mnist/mlp_custom_mnist.pt")
net.load_state_dict(state_dict)
train.test_model(net, testloader, loss_fn=loss_fn)
mask = magnitude_pruning(net, pruning_rate=0.25)
print("Number of ones in mask:", number_of_ones_in_mask(mask), "\n")
train.test_model(net, testloader, loss_fn=loss_fn)<jupyter_output>TESTING - loss 7674.807939052582 - performance 0.9612<jupyter_text>We see that there's a drop in performance. This is a characteristic of pruning. Rarely a simple application of pruning leaves the performance untouched or better. Due to that, we must procede with a re-training of the ANN.<jupyter_code>def train_epoch(model, dataloader, loss_fn, optimizer, loss_meter, performance_meter, performance, device, mask, layers_to_prune=["1", "4", "7", "10"]):
    for X, y in dataloader:
        X = X.to(device)
        y = y.to(device)

        optimizer.zero_grad() 

        y_hat = model(X)

        loss = loss_fn(y_hat, y)

        loss.backward()

        ##### we must neutralize the gradient on the pruned params before the optimizer takes a step ####

        if mask is not None:
            for (name, param), m in zip(model.named_parameters(), mask):
                if any([l in name for l in layers_to_prune]):
                    param.grad *= m

        ######

        optimizer.step()

        acc = performance(y_hat, y)

        loss_meter.update(val=loss.item(), n=X.shape[0])
        performance_meter.update(val=acc, n=X.shape[0])
def train_model(model, dataloader, loss_fn, optimizer, num_epochs, checkpoint_loc=None, checkpoint_name="checkpoint.pt", performance=accuracy, lr_scheduler=None, device=None, mask=None, params_type_to_prune=["weight", "bias"]):
    if checkpoint_loc is not None:
        os.makedirs(checkpoint_loc, exist_ok=True)

    if device is None:
        device = use_gpu_if_possible()
    
    model = model.to(device)
    model.train()

    for epoch in range(num_epochs):
        loss_meter = AverageMeter()
        performance_meter = AverageMeter()

        print(f"Epoch {epoch+1} --- learning rate {optimizer.param_groups[0]['lr']:.5f}")

        train_epoch(model, dataloader, loss_fn, optimizer, loss_meter, performance_meter, performance, device, mask)

        print(f"Epoch {epoch+1} completed. Loss - total: {loss_meter.sum} - average: {loss_meter.avg}; Performance: {performance_meter.avg}")

        if checkpoint_name is not None and checkpoint_loc is not None:
            checkpoint_dict = {
                "parameters": model.state_dict(),
                "optimizer": optimizer.state_dict(),
                "epoch": epoch
            }
            torch.save(checkpoint_dict, os.path.join(checkpoint_loc, checkpoint_name))
        
        if lr_scheduler is not None:
            lr_scheduler.step()

    return loss_meter.sum, performance_meter.avg
num_epochs = 10
lr = 0.01
optimizer = torch.optim.Adam(net.parameters(), lr = lr)

train_model(net, trainloader, loss_fn, optimizer, num_epochs, device="cpu", mask=mask)
train.test_model(net, testloader, loss_fn=loss_fn)
mask = magnitude_pruning(net, pruning_rate=0.25)
print("Number of ones in mask:", number_of_ones_in_mask(mask), "\n")
train.test_model(net, testloader, loss_fn=loss_fn)<jupyter_output>TESTING - loss 3899.6647787094116 - performance 0.9791<jupyter_text>We can see that the number of ones in the mask is the same as before! So we couldn't prune another 25% of the parameters.  Iterative Magnitude Pruning (IMP)To be able to prune again the parameters in the network, we compute the threshold removing the parameters that are already zero, so that we compute the pruning rate only on the non-zero parameters.<jupyter_code>def magnitude_pruning(model, pruning_rate, layers_to_prune=["1", "4", "7", "10"]):
    # 1. vectorize distribution of abs(parameter)
    '''
    model.named_parameters returns a tuple:
      element 0 is the name of the parameter
      element 1 is the parameter itself
    any(<list of booleans>)
    checks whether any of the conditions inside the list is True.
    [l in pars[0] for l in layers_to_prune] checks whether "1", "4", "7", and "10" are inside the parameter name
    hence, there will be up to 1 True in the corresponding list: in that case, any(...) returns True
    '''
    params_to_prune = [pars[1] for pars in model.named_parameters() if any([l in pars[0] for l in layers_to_prune])]
    flat = torch.cat([pars.abs().flatten() for pars in params_to_prune], dim=0)
    print("Parameters dimensions: ", end="")
    print(flat.shape)

    # 2. sort this distribution
    flat = flat.sort()[0]

    # 2.5. remove the zeros, thus the parameters already pruned
    flat = flat[flat.nonzero().detach()]
    print("Pruned parameters dimensions: ", end="")
    print(flat.shape)

    # 3. obtain the threshold
    position = int(pruning_rate * flat.shape[0])
    thresh = flat[position]

    # 4. binarize the parameters & 5. compose these booleans into the mask
    # 6. obtain the new structure of parameters
    '''
    I do this process with a for cycle instead of a list comprehension for clarity
    * if the layer is a layer to prune → populate the mask with 1s and 0s
    * otherwise → just populate the mask with ones
    By doing so, I can immediately apply the mask to the model as well...
    '''
    mask = []
    for pars in model.named_parameters():
        if any([l in pars[0] for l in layers_to_prune]):
            m = torch.where(pars[1].abs() >= thresh, 1, 0)
            mask.append(m)
            pars[1].data *= m
        else:
            mask.append(torch.ones_like(pars[1]))
            # no need to multiply as it's all 1s

    # 7. what do we need to return?
    return mask
# load pretrained model
state_dict = torch.load("../labs/models_push/mlp_custom_mnist/mlp_custom_mnist.pt")
net.load_state_dict(state_dict)

train.test_model(net, testloader, loss_fn=loss_fn)
for i in range(10):
    print(f"Iteration {i+1}.\n")
    mask = magnitude_pruning(net, pruning_rate=0.25)
    print("Number of ones in mask:", number_of_ones_in_mask(mask), "\n")

    optimizer = torch.optim.Adam(net.parameters(), lr = lr)
    train_model(net, trainloader, loss_fn, optimizer, num_epochs, device="cpu", mask=mask)
    train.test_model(net, testloader, loss_fn=loss_fn)
    print("\n")<jupyter_output>Iteration 1.

Parameters dimensions: torch.Size([15866])
Pruned parameters dimensions: torch.Size([15866, 1])
Number of ones in mask: 0.7535114978247358 

Epoch 1 --- learning rate 0.01000
Epoch 1 completed. Loss - total: 12903.265724182129 - average: 0.21505442873636882; Performance: 0.93435
Epoch 2 --- learning rate 0.01000
Epoch 2 completed. Loss - total: 9346.30019569397 - average: 0.15577166992823283; Performance: 0.9516166666666667
Epoch 3 --- learning rate 0.01000
Epoch 3 completed. Loss - total: 7972.33712720871 - average: 0.13287228545347848; Performance: 0.9576666666666667
Epoch 4 --- learning rate 0.01000
Epoch 4 completed. Loss - total: 7332.6753578186035 - average: 0.12221125596364339; Performance: 0.9619166666666666
Epoch 5 --- learning rate 0.01000
Epoch 5 completed. Loss - total: 6764.026443362236 - average: 0.11273377405603727; Performance: 0.9646
Epoch 6 --- learning rate 0.01000
Epoch 6 completed. Loss - total: 6354.145076751709 - average: 0.10590241794586182; Perfor[...]
{'fraction_non_alphanumeric': 0.10813317258647871, 'fraction_numerical': 0.04785310210516554, 'mean_word_length': 3.536811106436685, 'pattern_counts': {'":': 12, '<': 12, '<?xml version=': 0, '>': 14, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 3, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'chain_length': '3', 'dir': 'jupyter-structured-clean-dedup', 'id': '345501', 'license': 'MIT', 'n_tokens_mistral': 3798, 'n_tokens_neox': 3393, 'n_words': 1142, 'path': 'homeworks/Homework_05.ipynb', 'repo_name': 'AngieC96/DSSC_DL_2021', 'size': '10783'}