Tensor and nn.Module Pruning

[mickypaganini]

🚀 Tensor and nn.Module Pruning

Tensor method and/or nn util to sparsify tensors and/or model, according to various pruning techniques in the literature.

Motivation

State-of-the-art deep learning techniques rely on over-parametrized models that are hard to deploy. On the contrary, biological neural networks are known to use efficient sparse connectivity. It's important to identify best techniques to compress models by reducing the number of parameters in them, in order to reduce memory, battery, and hardware consumption without sacrificing accuracy, deploy lightweight models on device, and guarantee privacy with private on-device computation. On the research front, pruning is used to investigate the differences in learning dynamics of over-parametrized and under-parametrized networks, to study the role of lucky sparse subnetworks and initializations ("lottery tickets" [1]), as a destructive neural architecture search technique, and others.
Goal of this feature: harmonizing pruning practices by providing a standard interface in PyTorch.
Target audience: researchers, engineering and product teams.

Pitch

Minimalist API, with deeper flexibility for power-users.

At the tensor level, this could look as follows:

t = torch.randn(3, 2, 4)
# e.g.: randomly mask 6 entries
pruned_t = t.prune(method='random', amount=6)
# e.g.: prune bottom 80% of entries by absolute value
pruned_t = t.prune(method='L1', amount=0.8)
# e.g.: prune 50% of channels along the last dimension by L1 norm
pruned_t = t.prune(method='L1Structured', amount=0.5)
# e.g.: prune 2 channels along the 0th dimension by L2 norm
pruned_t = t.prune(method='L2Structured', amount=2, axis=0)
# e.g.: prune 1 channel along the last dimension by L0 norm
pruned_t = t.prune(method='LnStructured', n=0, amount=1)

A not-in-place .prune method will return a torch.Tensor of the same type and size as the one it acts on.
In-place pruning supported via t.prune_(...).

At the model level, this will require a bit of thinking but should follow similar API patterns. This is important because not all pruning methods make sense on all parameters in a model (pruning conv kernels != pruning biases != pruning RNNs != pruning in the presence of batch norm, etc.).
First, we should have a sensible, well-documented default behavior for the average-user's API, where a call to net.prune(method='L1', amount=0.8) defaults to pruning PyTorch "prepackaged" modules (such as linear, conv, and recurrent layers) in some sensible, expected way.
Most power users though would probably want to prune custom layers, or prune different layer types or layers at different depths using different pruning methods or pruning method parameters. This could be specified via a dictionary, which maps parameter names (contained in net.state_dict().keys()) to a pruning method and its parameters:

{
    'features.0.weight' : L1PruningMethod(amount=0.8),
    ...
}

Similar to the tensor operations, model-level pruning could return a copy of the model, or act on the model in place.

Pruning methods can be used during or post- training; this implementation will be training-loop-agnostic: the user will have to take care of writing their own training loop to decide when to prune and what to do with the pruned object (re-initialize and retrain, finetune, etc.).

Alternatives

Depending on where this will live within the codebase, t.prune(method=method, **kwargs) could also look like: torch.nn.utils.pruning_function(t, **kwargs) or torch.nn.utils.pruning_function(module=module, name=param_name, **kwargs). I personally would prefer the first option because the kwargs are parameters of the pruning method itself, while t is the tensor it acts on (some pruning methods will also have to take in some data X or (X, y) when they're applied), but the last option is more in line with how, say, weight_norm is implemented. Perhaps, for Module-level application, following the example of the weight_norm implementation using hooks will make this more PyTorch-y, but I don't know if we want to sacrifice the ability to act directly on a tensor that is not part of a Module. Would that go into torch.nn.functional? Open to suggestions here.

cc @albanD @mruberry @jbschlosser

[ezyang]

cc @soumith. This sounds generally reasonable, but I am a little wary about signing up to figure out what "reasonable default pruning" behavior is for every module. Maybe this could live out of tree for the short term, and we can assess as people get more experience with what is 100%, obviously useful functionality, and what is a bit more opinion based and changes as research evolves?

[ezyang]

Some other comments that I had in private conversation with @mickypaganini:

• There's probably some API bikeshedding to do on the exact signature of prune (and whether or not it lives as a method or function). All of these possibilities are easy to implement, we just have to decide what to do. One non-obvious sticking point is that writing binding code for operators which take "enums" (like the string L1/L1Structured) is a bit more involved, and so we should also consider just having separate functions/methods for each pruning mechanism (unless there are other design reasons why having it as an enum is a good idea; for example, there are other functions that need to use the same enum, or if users will want to parametrize over pruning strategy uniformly over many calls to the prune method)
• If you do need to implement some actual CPU/CUDA kernels, it is probably easier to contribute them straight to PyTorch repository
• For a specialized topic like this, it is very helpful to have someone who steps up to do continual maintenance and bugfixes for it
• PyTorch prefers to put "obviously good ideas" in the library, and leave people the space to experiment

More motivation for merging directly to PyTorch, from Michela:

The only issue with [an external library for pruning] is that it'll become the (N+1)th solution for pruning models. I don't think anybody needs that. What we need is one centralized canonical way of doing it

Some other reference material:

• https://medium.com/tensorflow/tensorflow-model-optimization-toolkit-pruning-api-42cac9157a6a

[mickypaganini]

Thanks for summarizing.
Main road block now: deciding where tensor pruning will live.

1. If we go the nn.utils route: will the signature of nn.utils.mypruningfunction be nn.utils.mypruningfunction(module, name, **kwargs) as it's done in the weight_norm example I linked? That specific .apply method uses hooks, so it assumes that the tensor I want to prune is inside some module which will undergo some forward pass. It won't support independent tensors being pruned outside of this world.
2. To counter this, should we have a separate tensor method? One should be able to prune any given tensor even if they are not part of a module.

It feels weird to have the thing that controls pruning application in the forward and backward pass also handle the actual pruning logic implementation. Where should I have that logic live?

• I'm happy to make each pruning technique a separate function/method and avoid string identifiers.
• I'm also happy to maintain these functionalities long term.
• We can table default module pruning for each specific nn.Module kind for a later discussion.

[michaelklachko]

FYI, there's also temporal (lifetime) sparsity for activations:
https://arxiv.org/abs/1409.2752
https://arxiv.org/abs/1903.11257

[michaelklachko]

Also, typically activation pruning is more aggressive during training (might even be disabled during test in some cases).

[mickypaganini]

I started prototyping this idea using simple pruning methods: random and L1-based unstructured pruning, and random and Ln-based structured pruning. These are simple methods that are either random or only depend on the magnitude of the weights. You can check it out on my fork.
[Thanks @michaelklachko for pointers to other pruning method. We have a long list of candidate methods to implement in future iterations, including activation-based methods, weight evolution methods, etc.]

Pruning is currently located under torch/nn/utils, as suggested. The implementation design follows the example of weight_norm.py, using forward pre hooks to reparametrize a parameter in terms of the computed mask and the original parameter's value.

One would interact with pruning in a similar way as with weight norm, i.e. through functions like ln_structured_pruning(module, name, amount, n, axis), that applies pruning, and remove_pruning(module, name) that removes the hook and the reparametrization, and permanently substitutes module[name] with the pruned version of the parameter.

Pruning methods can be composed using a PruningContainer. This enables iterative pruning.

I'm still writing tests to ensure correctness.
Looking forward to your feedback!

[michaelklachko]

Note that pruning operation is similar to quantization. So perhaps the two can be implemented in a similar manner, or even combined. For example, binarizing ReLU activations, and ternarizing weights can be considered a form of pruning.

[michaelklachko]

Here's one example of how quantization op can be integrated into a layer. In my opinion, it makes sense to implement weight pruning/quantization a wrapper for a layer (e.g. like weightnorm), however for activations it's probably better to have it as a separate layer type (e.g. like batchnorm).

[soumith]

@mickypaganini thanks a ton for the proposal, and your WIP branch.

I reviewed the overall (new) design using hooks, and focused more around nn.Module rather than tensors themselves. I think this design is less invasive and makes a lot more sense -- because pruning is primarily targeted for weights anyways. It looks great.

I have some bikeshedding points. I think you'd want to place the code in nn.utils.prune and then remove pruning from all of the functions that you implemented.
For example:

# Current
import torch.nn.utils as utils

utils.ln_structured_pruning(...)
utils.remove_pruning(...)

# New
import torch.nn.utils.prune as prune

prune.ln_structured(...)
prune.remove(....)

For each doc snippet, you probably also want to add Example: apart from Args: and Output:
I'm excited to see this unlock a bunch of pruning research.

@michaelklachko I disagree that quantization should be combined with pruning because:

• Pruning is experimental in it's research life-cycle, not many think about optimized code for a pruned model for all of the pruning methods that folks are currently doing research on.
• Quantization on the other hand is extremely mature from the scientific computing perspective and in the neural network land -- there are optimized code and algorithms implemented for scale and shift quantization.
• So, treating quantization as a subset of pruning in the interface and code-paths will either limit the interface we expand pruning to, because we want it to also run fast OR will limit the quantization to be of limited performance.
  So, I think I prefer quantization to be a separate, fully-fleshed out interface, as is being implemented in Model Quantization for PyTorch (Proposal) #18318 instead of being conflated with more exotic pruning techniques.

[soumith]

@mickypaganini one thing to keep in mind though is that using forward hooks has it's limitations.
For example, if you want to prune based on post-activation information, it's not very straight-forward to do so.

Example:

import torch.nn as nn
import torch.nn.functional as F
import torch.nn.utils.prune as P

class Net(nn.Module):
    def __init__(self):
        super(Net, self).__init__()
        self.conv1 = P. random_structured(nn.Conv2d(1, 20, 5, 1), 'weight', 0.5, 1)

    def forward(self, x):
        y = F.relu(self.conv1(x))
        z = F.max_pool2d(y, 2, 2)
        return F.log_softmax(z, dim=1)

Here, you cannot prune based on the values of y using the current interface, you can prune only based on the values of x, conv1.weight and conv1.bias and at best conv1's output. y is the value of the output of conv1 after a relu, which isn't materialized after all pre and post hooks of self.conv1 are already executed.
Is this okay, or do you see any pruning methods that cannot be done because of this limitation in the interface?