[proposal] [discussion] Refactor pruning/weight_norm using new Reparametrization functionality + actually deprecate old impl of SpectralNorm

[vadimkantorov]

Currently the weight_norm and spectral_norm are patching a passed module + implement special functions for adding/removing these from a module.

Some ideas for refactoring to make it less tricky:

• provide a stable signature for getting weight, then they can be cleanly used with methods such as torch.matmul and F.conv2d
• if module patching (adding some new buffers as parameters and registering a hook) is needed and is a reasonable pattern, provide a user-facing stable abstraction for it (especially adding and removal of parameters). It seems we have a chain of decorators-hooks, and each of them may have some extra buffers, and currently they are all patched into the passed module object.

cc @jerryzh168 @jianyuh @dzhulgakov

[fmassa]

Could you illustrate your proposal with some code snippets? It will be easier to understand

[vadimkantorov]

For first point, spectral norm has compute_weight method, but it’s defined in an internal structure. I think a global method like constrain_weight_with_spectral_norm would be more discoverable.

For second point, to come up with a good abstraction, one needs to know why the current module object patching was chosen vs defining a new wrapping module. I don’t know the answer.

@t-vi may have some ideas, he wondered about a better abstraction also in #7298 (comment)

[zou3519]

We're trying to figure out what this would look like -- code snippets would definitely help

[t-vi]

So what it looks like to me as a user is that I would want to have "computed parameters".

• They have raw parameters as inputs.
• They are updated when the inputs change.
• They are cached after they are updated and can be used just like other tensors.

I would propose to

• Have a ParameterModel class that actually is a Module subclass and is initialized with the name(s) of the parameter(?) that is computes. It can have parameters just like a model. It's forward computes the parameters. It should optionally also have a way to set parameters (i.e. doing the inverse transform), e.g. for initialization.
• Have a "parameters were updated hook" that is called by the optimizer which invalidates the computed parameters, e.g. by setting them to None or so.
• Have a thing in Module.getitem that calls the ParameterModel instance when a computed parameter is None and is accessed.

This does not solve how to call the hook when someone else overwrites the parameter, but hey. I'm not set on any of the specifics. I can whip up a prototype if that helps.

So the main functional difference to how spectral normalization is currently implemented is that it replaces the forward_pre hook with an update hook. The advantage to the user is that she can see this as a modification of parameters rather than the modules having the parameters.

[t-vi]

So to put some code here (not remotely tested enough):
https://github.com/t-vi/pytorch/tree/calculated_param
has the adaptation of Module to introduce CalculatedParameter (and Parameter, optim.SGD, optim.Adam are adapted).
The biggest weakness (in my opinion) of the implementation is that it "needs help" to detect parameter updates. The typical case will be the optimizer doing it, so the optimizer needs to call param.updated().

Note how the spectral norm class now looks like a regular module - to the point where you would not need to have the utility spectral_norm and remove_spectral_norm functions (the latter could just be replaced m.weight = nn.Parameter(m.weight.data) if you don't need to programmatically control the name and check whether weight was a parameter).

While it does add some trickery in module.py, I think the user of calculated parameters gets a much better interface.

[t-vi]

An alternative to trying to chase parameter changes might be to use the pre_forward_hook within module. It would seem to depend on the application whether you want something giving a new result at every invocation even when the inputs didn't change or whether that is wasteful. This could be mitigated by having the update only be done in training mode or some such. One could also simplify the logic between the module and the CalculatedParam by keeping a calculation counter and check for None or the calculation counter to indicate out of date instead of having a hook that sets the parameter to None.

[ngimel]

Incidentally, we have also been working on abstracting reparameterization https://github.com/NVIDIA/apex/tree/master/apex/reparameterization. Current implementation puts boilerplate code to reparameterization.py, with weight norm defining just particular weight recomputation functions. The recomputation of weight is controlled by self.evaluated which is set to true at evaluation and toggled to false in backward (which may be crude enough, but worked not to recompute weight at each time step of rnn).
It is still a hook-based implementation, and conceptually I agree that all the same things can be achieved by module wrappers.

[t-vi]

@ngimel Nice, that looks like a clean variant of how weight_norm is currently done in master.

One thing I am wondering (and I think that is the main difference to my proposed implementation strategy): Is this something that works at module or at parameter level. I see it more naturally at the parameter level, while the other solutions seem to look more at the module level.

One thing that isn't easily achieved with my demo implementation is having multiple parameters that need to be calculated jointly. This could be achieved by returning a dict and have getattr check for that.

In terms of usability ("less boilerplate code") I think that allowing the parameter computation to be in a nn.Module subclass is a more natural solution than wrappers.

[ngimel]

@t-vi I've asked the author of that implementation to come here to give his perspective. CalculatedParameters in nn.Module that you have look nice, but they do require modifications to the core and optimizers, whereas apex approach is more self-contained. As you note, joint recalculation of multiple parameters does not come naturally, neither to apex nor to CalculatedParameters. Also, I think all current approaches break parameter sharing, and this is something that has to be figured out going forward.

[t-vi]

Actually, I could move the caching to the __call__ of CalculatedParameter. That would solve reduce the changes to Module and enable shared parameters.

(Edit: I did now, and also provided a tentative interface for multiple parameters by returning a dict from forward and the caching matching the suffix - not 100% clean, but it might be a start.)

[vadimkantorov]

Do we want parameter calculation to be able to provide gradients by itself?

[t-vi]

If you want to manipulate gradients, you probably want to do that using Function or a backward hook.

[raulpuric]

Chiming in on this discussion. I'm the author of the reparameterization @ngimel shared earlier.
I've actually thought of 3 of these solutions in the past, and I'll be providing my thoughts on them in addition to the thoughts on my implementation.

My/hooking implementation:
Pros:

• Abstract interface for reparameterization
• wrapper to add hooks to all parameters in a module recursively
• handles some degree of parameter sharing.
• Works behind the scenes. No modifications to module aside from adding new parameters. The module can be used in the exact same way as before.

Cons:

• can be hard/hacky to use with parameter sharing.
• saving/reloading model is not user-friendly. One can't load the weight_norm (g/v) parameters into a model without weight norm.

Optimizer handles reparameterization:
Pros:

• Saving/reloading model is user friendly as it does not modify anything in the model
  • works behind the scenes disentangled from the module
• Offloading responsibility for doing reparameterizations makes sense in the context of weight norm, as one of the goals is to change how the parameters are optimized.

Cons:

• Unusual optimizer usage is hard to allow (skipping optimization if gradients/loss blowup)

Calculated/Lazy Parameter:
Personally I'm a fan of this option, and had thought about it in the past, but didn't have time to implement it in a good fashion. If someone can implement a general class for this in a very safe PyTorch-y way, that also makes saving/reloading parameters easy, I would be in support of this option.
Pros:

• Works at the parameter level and does not add new parameters to the module
  • Easy to swap param with a non-reparameterized version of the parameter (helps with save/load)
  • should be able to handle weird optimization strategies so long as .data and .grad are managed well

Cons:

• Works at the parameter level so managing these might be hard.
  • Dictionary is a good first step, but needs something more. A hook might be a good solution here, since no new parameters/functionality are getting added to the module itself

Module wrapper:
Pros:

• It can handle everything that all the other options can (fprop, bprop, param updates, etc)

Cons:

• Module wrappers can complicate code. In order to access attributes/functionality of the wrapped module you have to add an additional .module everywhere

The ideal sort of setup I envision is either one module wrapper class that handles everything for the user, or a combination of CalculatedParameters and hooks/module wrapper. (Preferably the latter)

I think though in general PyTorch needs some sort of module wrapper nn.Module subclass that absorbs most if not all the functionality/attributes of the wrapped module instance. This would alleviate a lot of code considerations that have to be done when using things like (Distributed)DataParallel which forces us to add model.module all over our code.

[t-vi]

@raulpuric Thank you for commenting. If you have a moment to look at it, what do you think are the key things missing from the calculated param prototype ( https://github.com/t-vi/pytorch/tree/calculated_param ) relative to your vision?

[vadimkantorov]

Imagine we have a parameter precalc with for an op (pseudo-code):

weight, bias, other_param = ..., ..., ....
weight, bias = apply_spectral_norm_or_something_else(weight, other_param)
y = F.conv2d(x, weight, bias)

1. Are we supporting this functional scenario easily? (I think that we should)
2. Shoud we support differentiating to other_param? (I think the user should have a choice. The hooks-based solution doesn't allow it, does it? And the CalculatedParameter?)
3. Do we want to support this kind of joint param pre-calculation?

For the question of wrapper vs patching, I think wrapper is still cleaner (even if the user has to type .module to get to original module and may abstruct model inspection a bit), because for the case of chained parameter-precalc (imagine applying weight-norm and spectral-norm together, if it made sense) the order of the precalc ops is clear to user.