Add pruning tutorial

[mickypaganini]

Add tutorial for pruning functionalities discussed in feature request issue pytorch/pytorch#20402

Depends on: pytorch/pytorch#24076

[netlify]

Deploy preview for pytorch-tutorials-preview ready!

Built with commit 261d093

https://deploy-preview-605--pytorch-tutorials-preview.netlify.com

[SethHWeidman]

@mickypaganini, tutorial looks great. It wasn't building, so I made a tiny change and now it is. @brianjo, is this because @mickypaganini doesn't have write access to the repo? I think that's the issue I was having before.

cc @jlin27 FYI

[brianjo]

In the build (and when I test this locally), I get the error ModuleNotFoundError: No module named 'torch.nn.utils.prune'. We build on 1.2 release. Update - prune hasn't been merged into core yet, so this is going to be on hold for a bit. Thanks!

[mickypaganini]

Thanks for taking a look so quickly, @SethHWeidman @brianjo.

Yep, as Brian said and as I mentioned above, this depends on pruning actually getting merged into PyTorch (PR here: pytorch/pytorch#24076). Until then, it won't build. I will comment here once I have updates on that.

[mickypaganini]

@SethHWeidman @brianjo if you build off of pytorch master, this tutorial should build fine now.
pytorch/pytorch#24076 was successfully merged 🎉

[cranmer]

Hello... I found this tutorial via google search. I think it will be nice to include.

It's not immediately clear to me when the pruning is just an effective pruning implemented with masking and when it leads to a network that actually has a smaller memory footprint.
It seems that the part of the tutorial that uses prune.remove makes the pruning permanent. Should I see smaller memory footprint at that time? Or do I need to serialize it and read then load the model?

Sorry if these are naive questions, but hopefully they are helpful in reaching the audience the tutorial is aimed for.

[mickypaganini]

Hi @cranmer,
Hopefully this tutorial will be included soon (cc: @soumith).

As is, this module is not intended (by itself) to help you with memory savings. All that pruning does is to replace some entries with zeroes. This itself doesn't buy you anything, unless you represent the sparse tensor in a smarter way (which this module itself doesn't handle for you). You can, however, rely on torch.sparse and other functionalities there to help you with that. To give you a concrete example:

import torch
import torch.nn.utils.prune as prune

t = torch.randn(100, 100)
torch.save(t, 'full.pth')

p = prune.L1Unstructured(amount=0.9)
pruned = p.prune(t)
torch.save(pruned, 'pruned.pth')

sparsified = pruned.to_sparse()
torch.save(sparsified, 'sparsified.pth')

When I ls, these are the sizes on disk:

21K sparsified.pth
40K pruned.pth
40K full.pth

By the way, before calling prune.remove, you can expect you memory footprint to be a lot higher than what you started out with, because for each pruned parameter you now have: the original parameter, the mask, and the pruned version of the tensor. Calling prune.remove brings you back to only having a single (now pruned) tensor per pruned parameter. Still, if you don't represent these pruned parameters smartly, the memory footprint at this point won't be any lower than what you started out with.

[takshshila]

The tutorial mentioned pruning only a convolutional layer. Is there any method to prune the complete
BERT model?

[Christian-lyc]

Yes, I'm also curious about how to apply it into the complete model.

[takshshila]

Yes, I'm also curious about how to apply it into the complete model.

Hi used import torch.nn.utils.prune as prune to prune attention weights of the BERT model. e.g.

`parameters_to_prune = ()
layers = len(bert_model.bert.encoder.layer)

for i in range(layers):
    parameters_to_prune += (
        (bert_model.bert.encoder.layer[i].attention.self.key, 'weight'),
        (bert_model.bert.encoder.layer[i].attention.self.query, 'weight'),
        (bert_model.bert.encoder.layer[i].attention.self.value, 'weight'),
    )

prune.global_unstructured(
    parameters_to_prune,
    pruning_method=prune.L1Unstructured,
    amount=pamount,
)`