Pruning doesn't affect speed nor memory usage

[Cyril9227]

Hello,

I'm experimenting with the newly pruning feature on Pytorch and going thru the tutorial. My issue is that after pruning weights of a ResNet, the inference time and the memory footprint is exactly the same.

How do I get speed-up and reduced memory footprint using pruning in Pytorch ?

My code for pruning is pretty simple so far :

for name, module in model.named_modules():
    # prune 20% of connections in all 2D-conv layers
    if isinstance(module, torch.nn.Conv2d):
        prune.l1_unstructured(module, name='weight', amount=0.2)
        prune.remove(module, 'weight') # make it permanent
    # prune 40% of connections in all linear layers
    elif isinstance(module, torch.nn.Linear):
        prune.l1_unstructured(module, name='weight', amount=0.4)
        prune.remove(module, 'weight')

For timing I use :

    start = time.time()
    # get predictions
    output = model(inputs)
    end = time.time()  
    print(end - start)

Any help is welcome,
thanks

[gchanan]

Since this is a question, you might get more help on the forums (https://discuss.pytorch.org/).

CC @mickypaganini.

[mickypaganini]

Correct, it doesn't. See the conversation here as well: pytorch/tutorials#605 (comment)

[Cyril9227]

Correct, it doesn't. See the conversation here as well: pytorch/tutorials#605 (comment)

Thank's for the answer @mickypaganini , this is what I thought.

Any idea of the easiest way to sparsify the model once pruned ?

Thank's

[zeng-hello-world]

The pytorch pruner just mask zeros out in weight tensor, so number of the computation ops is still the same.

[mickypaganini]

As far as I know, for memory concerns, you can use .to_sparse().
Example:

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

t = torch.randn(1000, 1000)
torch.save(t, 'original.pth')

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

sparse_t = pruned_t.to_sparse()
torch.save(sparse_t, 'sparse.pth')

Then ls -lht *.pth should return:

196K sparse.pth
3.9M pruned.pth
3.9M original.pth

This only works if you prune your tensor by a large enough amount that representing it in coordinate space is actually memory efficient.

You can look at torch.sparse to see what ops are then supported on these sparse tensors. Note that this is marked as experimental and the API might change in the future: https://pytorch.org/docs/stable/sparse.html

I'm not sure about the computational speed-up concerns... @bwasti, maybe?

[albanD]

Corresponding PR has been merged.