[Python] How to perform global structured pruning in PyTorch

I am training on CIFAR10 the following simple CNN

class SimpleCNN(nn.Module):
def __init__(self):
super().__init__()
self.conv1 = nn.Conv2d(3, 32, kernel_size=3, padding=1)
self.conv2 = nn.Conv2d(32, 64, kernel_size=3, padding=1)
self.pool = nn.MaxPool2d(2, 2)
self.relu = nn.ReLU()

self.fc1 = nn.Linear(64 * 8 * 8, 128, bias=False)
self.fc2 = nn.Linear(128, 128, bias=False)
self.fc3 = nn.Linear(128, 10, bias=False)

def forward(self, x):
x = self.pool(self.relu(self.conv1(x)))
x = self.pool(self.relu(self.conv2(x)))
x = x.view(-1, 64 * 8 * 8)
x = self.relu(self.fc1(x))
x = self.relu(self.fc2(x))
x = self.fc3(x)
return x


I would like to perform global structured pruning on the last part of the network, the MLP part (so fc1, fc2, fc3) with a percentage cutoff. Basically I would like to cut off x percent of the neurons (and all the relative connections) on the basis of the total connectivity. In PyTorch there is a function that performs a similar job: ln_structured

import torch.nn.utils.prune as prune
prune.ln_structured(model.fc1, name='weight', amount=fraction_of_neurons_to_prune)
prune.ln_structured(model.fc2, name='weight', amount=fraction_of_neurons_to_prune)
prune.ln_structured(model.fc3, name='weight', amount=fraction_of_neurons_to_prune)


But the main problem is that this function will apply the fraction to prune layer by layer, and not globally; I want something to cutoff let's say 80 percent of the neurons in the MLP portion of the network, and not 80 percent layer by layer.

Is there a function that performs what I want? How could I write it myself? The absence of a function in PyTorch that performs this very common operation seems very strange to me, but I can't find anything.

Continue reading...