Save Pruned Model After Pruning

1

Pruning could also be used as a sort of neural architecture search method.
Say, I have a two-layer convolutional neural network.

conv1 = torch.nn.Conv2d(in_channels=3, out_channels=12)
conv2 = torch.nn.Conv2d(in_channels=12, out_channels=8)

If I used structured pruning and say the last three channels of the conv2 kernel was pruned to 0. Then the useful architecture of the pruned model should be

conv1 = torch.nn.Conv2d(in_channels=3, out_channels=9)
conv2 = torch.nn.Conv2d(in_channels=9, out_channels=8)

This architecture with the remaining parameters could be saved as a new model. The size of this new model should be smaller and the inference speed of the new model should be faster than the unpruned one.

I wonder if PyTorch is planning to add this interface or not.

1 Like
2

any updates on this?

3

If there aren’t any work on this, maybe we can create an extended class and create a push request?

4

@leimao Shouldn’t the ‘useful architecture’ post structured pruning be as:

conv1 = torch.nn.Conv2d(in_channels = 3, out_channels = 9)
conv2 = torch.nn.Conv2d(in_channels = 9, out_channels = 5)

because you are only pruning the last three channels/filters/kernel maps of conv2 layer? conv1 layer is untouched/unpruned.