Use forward_pre_hook to modify nn.Module parameters

Jiajun_Zha: <p>Every time before the forward() function of an nn.Module is called, I want to check and probably modify this nn.Module’s parameter, including weight and bias.</p>
<p>But the official source code of “register_forward_pre_hook” below doesn’t really say if this is achievable.</p>
<pre><code class="lang-auto">    def register_forward_pre_hook(self, hook):
        r"""Registers a forward pre-hook on the module.

        The hook will be called every time before :func:`forward` is invoked.
        It should have the following signature::

            hook(module, input) -&gt; None or modified input

        The hook can modify the input. User can either return a tuple or a
        single modified value in the hook. We will wrap the value into a tuple
        if a single value is returned(unless that value is already a tuple).

        Returns:
            :class:`torch.utils.hooks.RemovableHandle`:
                a handle that can be used to remove the added hook by calling
                ``handle.remove()``
        """
        handle = hooks.RemovableHandle(self._forward_pre_hooks)
        self._forward_pre_hooks[handle.id] = hook
        return handle
</code></pre>
<p>So the question is simple, could a forward pre hook modify the module?</p>
<p>P.S. I’m not really sure if this is the same question raised here <a href="https://github.com/pytorch/pytorch/issues/5790" class="inline-onebox" rel="noopener nofollow ugc">Add hookable weights · Issue #5790 · pytorch/pytorch · GitHub</a></p>

albanD: <p>Hi,</p>
<p>Yes, you can modify the Module in any way you want during that hook (the module is the first arg of the hook function).<br>
Note though that if you swap out weight Tensors, it might not behavior properly with other elements that were tracking the old Tensors like optimizers.</p>

sahilmn: <p>Do optimizers just use the name of the parameters to track 'em? I’m looking to use <code>register_forward_pre_hook</code> for equalized learning rate as described in the <a href="https://arxiv.org/abs/1710.10196" rel="noopener nofollow ugc">progressive GAN</a> paper. In simple terms, the weight is divided by a constant after every learning step.</p>
<p>In this case, if I simply divide the weights in <code>register_forward_pre_hook</code> by say <code>2</code>, how will it affect optimizer behavior?</p>

albanD: <p>The optimizers track the weights as Tensors directly.<br>
So as long as you modify that Tensor inplace, they will work just fine.</p>
<p>But if you do <code>model.weight = foo</code> then the new model.weight will not correspond to the one that is in the optimizer and so no learning will happen.</p>

iamanigeeit: <p>Hi albanD,</p>
<p>I want to mask some of the weights before <code>forward</code>, so the hook needs to take the mask as input. But the hook signature <code>hook(module, input)</code> doesn’t allow additional params.</p>
<p>Is what i want possible? Or should i just store the mask as a module buffer?</p>
<p>EDIT: i looked at the code in <code>torch.nn.utils.prune.py</code> and it looks like register_buffer is the right way.</p>

albanD: <p>Hi,</p>
<p>The hook can capture any variable from the parent scope. So if the mask is known ahead of time, you can do something like:</p>
<pre><code class="lang-python">def get_hook(mask):
  def hook(mod, input):
    # Do something with mod, input and mask
  return hook

mod.register_forward_pre_hook(get_hook(mask))
</code></pre>

iamanigeeit: <p>Thanks and just a note for others: if the <code>mask</code> is updated, you should delete the hook and register a new one with the updated mask.</p>
<pre><code class="lang-auto">mask = old_mask
handle = mod.register_forward_pre_hook(get_hook(mask))
mask = new_mask
handle.remove()
handle = mod.register_forward_pre_hook(get_hook(mask))
</code></pre>

msqiqi10: <p>Hi: albanD:</p>
<p>Thanks for the explaination. In my case, suppose I want to add a noise to the weights before the forward pass and restore the weights afterwards, what would be the better way for implementation without modifying the forward function of the module?<br>
My way of doing it is to add a pre hook and storing the noises into a dict:</p>
<pre><code class="lang-auto">def get_hook(mean=0.0, std=0.1):
    def hook(module, input):
        with torch.no_grad():
            if hasattr(module, 'weight'):
                noise = torch.randn_like(module.weight) * std + mean
                module.weight.add_(noise)
                noises[''] = noise
    return hook
</code></pre>
<p>and then delete the noise with forward hook:</p>
<pre><code class="lang-auto">def get_hook():
    def hook(module, input, output):
        with torch.no_grad():
            if hasattr(module, 'weight'):
                module.weight.add_(-noises[''])
    return hook
</code></pre>

albanD: <p>Hey!</p>
<p>This looks like a textbook use case for reparametrization! <a href="https://pytorch.org/tutorials/intermediate/parametrizations.html" class="inline-onebox">Parametrizations Tutorial — PyTorch Tutorials 2.2.2+cu121 documentation</a></p>