<jupyter_start><jupyter_text>Layer-wise Pruning RobustnessPrune each layer removing units iteratively according to the ranking of different attribution metrics. Plot the test loss and accuracy against the number of units that have been removed at each layer.<jupyter_code>%load_ext autoreload
%autoreload 2
import sys
sys.path.append("./..")
import torch
import torch.nn as nn
import numpy as np
import pickle
import matplotlib.pyplot as plt
from torchsummary import summary, torchsummary
from torchpruner import (Pruner, ShapleyAttributionMetric)
from experiments.utils import get_module_name, format_plt, map_method_vis, train, test
import experiments.models.cifar10 as cifar10

from torchpruner import (
    WeightNormAttributionMetric,
    RandomAttributionMetric,
    SensitivityAttributionMetric,
    TaylorAttributionMetric,
    APoZAttributionMetric,
    ShapleyAttributionMetric,
    find_best_module_for_attributions
)

# Fix seed for reproducibility
torch.manual_seed(1)
torch.backends.cudnn.deterministic = True
torch.backends.cudnn.benchmark = False
np.random.seed(1)

use_cuda = torch.cuda.is_available()
device = torch.device("cuda" if use_cuda else "cpu")
print (f"Using device: {device}")
# Get pretrained weights
!wget -O weights/CIFAR10-VGG16.pt "https://drive.google.com/u/0/uc?id=164ZfuLwocQRsnXPXKHr6Gbi5Pu7Ikihn&export=download"
train_set, test_set = cifar10.get_dataset()
# Using only 1000 samples from test set, just to speed up things a bit.
# Alternatively, this test could be run faster using a forward_partial function
test_set, _ = torch.utils.data.random_split(test_set, [1000, len(test_set)-1000])
test_loader = torch.utils.data.DataLoader(
    test_set, batch_size=1000, shuffle=False,
)
x_test, y_test = next(iter(test_loader))
x_test, y_test = x_test.to(device), y_test.to(device)

# Use part of the train data to compute attributions
train_set, _ = torch.utils.data.random_split(train_set, [1000, len(train_set)-1000])
val_loader = torch.utils.data.DataLoader(
    train_set, batch_size=100, shuffle=False,
)
    
loss = cifar10.loss
input_size = (3, 32, 32)

# Load pretrained model
def load_model():
    model, name = cifar10.get_vgg_model_with_name()
    # Load weights
    model.load_state_dict(torch.load("weights/CIFAR10-VGG16.pt"))
    model.to(device)
    model.eval()
    return model

# Display architecture and test performance
model = load_model()
summary(model, input_size=input_size, device=device.type)
original_loss, original_acc = test(model, device, loss, test_loader);
# Construct a function that returns the pruning graph for a VGG network
# The pruning graph is a list of tuples (module, [cascading_modules])
# where module is a prunable 'module' and cascading_modules contains all
# the modules that should be pruned in consequence of 'module' being pruned
def get_pruning_graph(vgg):
    modules = list(vgg.features.children()) + list(vgg.classifier.children())
    pruning = []
    current = None

    for module in modules:
        if any([isinstance(module, c) for c in [nn.Linear, nn.Conv2d]]):
            if current is not None:
                pruning[-1][1].append(module)
                pruning[-1][1].reverse()
            current = module
            pruning.append((module, []))
        elif (
            any([isinstance(module, c) for c in [nn.BatchNorm2d, nn.Dropout]])
            and current is not None
        ):
            pruning[-1][1].append(module)
    return pruning[::-1][1:]
# Initialize attribution metrics that will be used
methods = {
    "Weight Norm" : WeightNormAttributionMetric(model, val_loader, loss, device),
    "Random" : RandomAttributionMetric(model, val_loader, loss, device),
    "Sensitivity" : SensitivityAttributionMetric(model, val_loader, loss, device),
    "Taylor" : TaylorAttributionMetric(model, val_loader, loss, device),
    "Taylor signed" : TaylorAttributionMetric(model, val_loader, loss, device, signed=True),
    "APoZ" : APoZAttributionMetric(model, val_loader, loss, device),
    "SV" : ShapleyAttributionMetric(model, val_loader, loss, device),
    "SV mean+2std" : ShapleyAttributionMetric(model, val_loader, loss, device, 
                                             reduction=lambda x: np.mean(x, 0) + 2*np.std(x, 0))
}<jupyter_output><empty_output><jupyter_text>Run the test, for all layers and all attribution methods<jupyter_code>%%time
pruning_graph = get_pruning_graph(model)
print (f"Pruning {len(pruning_graph)} modules...")

log = {}
for module, cascading_modules in pruning_graph:
            
    module_name = get_module_name(model, module)
    log[module_name] = {}
    print (f"Running on {module_name}")
    
    for method_name, attribution in methods.items():
        log[module_name][method_name] = {"loss" : [], "acc" : []}
        pruned_indices = []
        runs = 3 if method_name in ["Random", "SV", "SV mean+2std"] else 1
        
        for r in range(runs):
            print (f"\t--> {method_name} (run {r+1})")
            log[module_name][method_name]["loss"].append([original_loss])
            log[module_name][method_name]["acc"].append([original_acc])

            # Compute attributions on a subset of the train data
            # `find_best_evaluation_module=True` makes sure attributions are computed
            # after BatchNorm and non-linearity
            scores = attribution.run(module, find_best_evaluation_module=True)
            ranking = np.argsort(scores)

            with torch.no_grad():
                # Start ablation procedure, measuring performance on the test data
                # We use simulated pruning (activations masked to zero)
                
                # 0. Find the next ReLU module, because pruning should happen after BN and non-linearity
                # Notice that VGG is composed by Conv -> BN -> ReLU blocks
                # During structured pruning, both Conv and BN are pruned, therefore
                # masking should happen at least after BN.
                module_nl = find_best_module_for_attributions(model, module)
                
                # 1. Compute activation of current module/layer.
                z = model.forward_partial(x_test, to_module=module_nl)

                # 2. Remove activations sequentially, starting with those having lowest attributions
                for i in ranking:
                    z.index_fill_(1, torch.tensor([i]).to(device), 0.)
                    y_pred = model.forward_partial(z, from_module=module_nl)
                    new_loss = loss(y_pred, y_test)
                    current_pred = y_pred.argmax(dim=1, keepdim=True)
                    new_acc = current_pred.eq(y_test.view_as(current_pred)).sum().float() / len(y_test)
                    log[module_name][method_name]["loss"][-1].append(new_loss.item())
                    log[module_name][method_name]["acc"][-1].append(new_acc.item())

# Dump all results to file for later analysis
pickle.dump(log, open("data/layerwise_pruning_results.p", "wb" ))<jupyter_output>Pruning 15 modules...
Running on classifier.4
	--> Weight Norm (run 1)
Activation detected: shifting evaluation after ReLU(inplace=True)
	--> Random (run 1)
Activation detected: shifting evaluation after ReLU(inplace=True)
	--> Random (run 2)
Activation detected: shifting evaluation after ReLU(inplace=True)
	--> Random (run 3)
Activation detected: shifting evaluation after ReLU(inplace=True)
	--> Sensitivity (run 1)
Activation detected: shifting evaluation after ReLU(inplace=True)
Activation detected: shifting evaluation after ReLU(inplace=True)
	--> Taylor (run 1)
Activation detected: shifting evaluation after ReLU(inplace=True)
Activation detected: shifting evaluation after ReLU(inplace=True)
	--> Taylor signed (run 1)
Activation detected: shifting evaluation after ReLU(inplace=True)
Activation detected: shifting evaluation after ReLU(inplace=True)
	--> APoZ (run 1)
Activation detected: shifting evaluation after ReLU(inplace=True)
Activation detected: shifting evaluation after ReLU(i[...]<jupyter_text>Plot the results<jupyter_code># Load the data from previously stored results
# and sort layers
plot_data = pickle.load(open("data/layerwise_pruning_results.p", "rb" ))
layers = plot_data.keys()
layers = sorted(layers, key=lambda x: -int(x.split(".")[1]), reverse=True)
layers = sorted(layers, key=lambda x: 0 if "classifier" in x else 1, reverse=True)

# Some general settings
plt.rcParams.update({'font.size': 11})

auc = {}
for idx, layer_id in enumerate(layers):
    layer_data = plot_data[layer_id]
    layer_name = f"{'Conv ' if 'features' in layer_id else 'Dense '} {idx}"
    fig, (ax_loss, ax_acc) = plt.subplots(1, 2, figsize=(10,3.0));
    
    for method_name, data in layer_data.items(): 
        label, color = map_method_vis(method_name)
        ticks = range(len(data["loss"][0]))
        loss_mean = np.mean(data["loss"], 0)
        loss_std = np.std(data["loss"], 0)
        acc_mean = np.mean(data["acc"], 0)
        acc_std = np.std(data["acc"], 0)
        
        # AUC
        auc_layer = np.sum(np.array(data["loss"]) - data["loss"][0][0], 1)
        n_units = len(data["loss"][0]) - 1
        if method_name not in auc:
            auc[method_name] = {"sum": auc_layer, "count": n_units}
        else:
            auc[method_name]["sum"] += auc_layer
            auc[method_name]["count"] += n_units
        
        # Plot with confidence intervals
        ax_loss.fill_between(ticks, loss_mean + loss_std, loss_mean - loss_std, alpha=.2, color=color);
        ax_loss.plot(ticks, loss_mean, label=label, color=color, lw=1);
        format_plt(ax_loss, f"{layer_name} - Loss", "# units removed", "Avg. Test Loss")

        ax_acc.fill_between(ticks, acc_mean + acc_std, acc_mean - acc_std, alpha=.2, color=color);
        ax_acc.plot(ticks, acc_mean, label=label, color=color, lw=1);
        format_plt(ax_acc, f"{layer_name} - Accuracy", "# units removed", "Test Accuracy (%)")
        # plt.savefig(f"cifar10_robustness_study_{layer_name}.pdf", bbox_inches='tight')
        

# Plot a single legend below
handles, labels = ax_loss.get_legend_handles_labels()
handles, labels  = zip(*sorted(zip(handles, labels), key=lambda t: 1 if "SV" in t[1] else -1))
fig = plt.figure(figsize=(10, 1))
fig.add_subplot(111)
plt.box(False)
plt.axis('off')
fig.legend(handles, labels,loc='upper center', ncol=4);
# plt.savefig(f"cifar10_robustness_study_legend.pdf", bbox_inches='tight')
# Print AUC
for method_name, auc_method in auc.items():
    s = auc_method["sum"]
    c = auc_method["count"] 
    print (f"{method_name} \n {np.mean(s/c, 0):.2f} +- {np.std(s/c, 0):.2f}")<jupyter_output>Weight Norm 
 0.47 +- 0.00
Random 
 0.48 +- 0.01
Sensitivity 
 0.47 +- 0.00
Taylor 
 0.47 +- 0.00
Taylor signed 
 0.64 +- 0.00
APoZ 
 0.56 +- 0.00
SV 
 0.35 +- 0.01
SV mean+2std 
 0.31 +- 0.00
{'fraction_non_alphanumeric': 0.11642315516922508, 'fraction_numerical': 0.017477344183465877, 'mean_word_length': 3.5194316757208526, 'pattern_counts': {'":': 2, '<': 11, '<?xml version=': 0, '>': 22, 'https://': 1, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 3, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'chain_length': '3', 'dir': 'jupyter-structured-clean-dedup', 'id': '655409', 'license': 'MIT', 'n_tokens_mistral': 3669, 'n_tokens_neox': 3426, 'n_words': 897, 'path': 'experiments/CIFAR-10 - VGG16 - Layerwise robustness.ipynb', 'repo_name': 'berserkr/TorchPruner', 'size': '10814'}