Pruning

Define any network, for instance:

class LeNet(nn.Module):
    def __init__(self):
        super(LeNet, self).__init__()
        # 1 input image channel, 6 output channels, 3x3 square conv kernel
        self.conv1 = nn.Conv2d(1, 6, 3)
        self.conv2 = nn.Conv2d(6, 16, 3)
        self.fc1 = nn.Linear(16 * 5 * 5, 120)  # 5x5 image dimension
        self.fc2 = nn.Linear(120, 84)
        self.fc3 = nn.Linear(84, 10)

    def forward(self, x):
        x = F.max_pool2d(F.relu(self.conv1(x)), (2, 2))
        x = F.max_pool2d(F.relu(self.conv2(x)), 2)
        x = x.view(-1, int(x.nelement() / x.shape[0]))
        x = F.relu(self.fc1(x))
        x = F.relu(self.fc2(x))
        x = self.fc3(x)
        return x

model = LeNet()

Now we will use Fluence's Pruner to create pruning processor.

from fluence.prune import Pruner

Now we will pass our model, which type of layers we want to prune and percentage

prune_proc = Pruner(model, ['conv', 'linear'], 0.675)

Now time to prune !!

prune_proc.perform_pruning('random')

You can define which type of pruning you want. To make pruning permanent,

prune_proc.make_permanent()

Get the pruned model

model = prune_proc.model