How to modify the module's inputs in computation graph after the forward

username: Zeyuan.Yin
cooked: <p>Hi, I am reproducing a research work about model pruning in generating adversarial examples. The paper link is: <a href="https://arxiv.org/abs/2208.08677" class="inline-onebox" rel="noopener nofollow ugc">[2208.08677] Enhancing Targeted Attack Transferability via Diversified Weight Pruning</a>. In Section 3.2, the description of model pruning is as follows,<br>
<div class="lightbox-wrapper"><a class="lightbox" href="https://discuss.pytorch.org/uploads/default/original/3X/4/a/4a8a2c4aefb5c737401435a1ed2fd2dbcff9c2e2.png" data-download-href="https://discuss.pytorch.org/uploads/default/4a8a2c4aefb5c737401435a1ed2fd2dbcff9c2e2" title="image"><img src="https://discuss.pytorch.org/uploads/default/optimized/3X/4/a/4a8a2c4aefb5c737401435a1ed2fd2dbcff9c2e2_2_690x102.png" alt="image" data-base62-sha1="aDpmLhJ5gTzwcNTew2dcx992a0q" width="690" height="102" srcset="https://discuss.pytorch.org/uploads/default/optimized/3X/4/a/4a8a2c4aefb5c737401435a1ed2fd2dbcff9c2e2_2_690x102.png, https://discuss.pytorch.org/uploads/default/original/3X/4/a/4a8a2c4aefb5c737401435a1ed2fd2dbcff9c2e2.png 1.5x, https://discuss.pytorch.org/uploads/default/original/3X/4/a/4a8a2c4aefb5c737401435a1ed2fd2dbcff9c2e2.png 2x" data-dominant-color="E7E7E7"><div class="meta"><svg class="fa d-icon d-icon-far-image svg-icon" aria-hidden="true"><use href="#far-image"></use></svg><span class="filename">image</span><span class="informations">1008×150 22.2 KB</span><svg class="fa d-icon d-icon-discourse-expand svg-icon" aria-hidden="true"><use href="#discourse-expand"></use></svg></div></a></div><br>
It means we use an original network for forward and its pruned network for backward to calculate the gradient.</p>
<p>The key to pruning the network before gradient computing in backpropagation is to modify the module’s inputs (containing the model’s weights) in the computation graph after the forward.</p>
<p>However, existing hook methods can not solve the problem,</p>
<ul>
<li><code>register_module_forward_pre_hook</code> can modify the input before the forward, but we want to do it after the forward.</li>
<li><code>register_module_forward_hook</code> can only modify the output instead of input after the forward.</li>
<li><code>backward_hook</code> can modify the gradient, but we want to modify the inputs (containing the model’s weights) in the computation graph</li>
</ul>
<p>I wonder whether there is any way to modify the input in the computation graph after the forward.</p>

username: soulitzer
cooked: <p>I’m not sure what you mean by use a pruned network for backward, but if you want to prune saved activations before using them to compute the backward pass, you could use <a href="https://pytorch.org/tutorials/intermediate/autograd_saved_tensors_hooks_tutorial.html" class="inline-onebox" rel="noopener nofollow ugc">Hooks for autograd saved tensors — PyTorch Tutorials 2.0.0+cu117 documentation</a>. With these hooks, the gradients would be computed as if the forward were computed using the pruned inputs. It should be noted that sometimes, the forward would save the output instead of input.</p>

username: Zeyuan.Yin
cooked: <p>Thanks a lot. Now I know I can use <code>unpack_hook</code> to get the save_tensors in the backward<br>
For my pruning task, for example, in a  linear layer,  y = W*x,  W and x are packed and saved into the save_tensors before the forward. I can use <code>unpack_hook</code> to unpack W and multiply W with a mask matrix to implement pruning to calculate the gradient.<br>
I am trying this idea and will post the code if it works</p>

username: Zeyuan.Yin
cooked: <p>The code is here, replace the <code>Conv2d</code> module with <code>Conv2dWithHooks</code></p>
<pre><code class="lang-auto">class Conv2dWithHooks(nn.Conv2d):
    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)

    def forward(self, input):
        with torch.autograd.graph.saved_tensors_hooks(self.pack, self.unpack):
            return super().forward(input)

    def pack(self, x): # nothing to modify in the forward pass
        return x

    def unpack(self, x): # modify weight before the backward pass
        if x.shape == self.weight.shape: # x is weight
            weight_mask = torch.ones_like(self.weight)   # example mask
            x = x * weight_mask
        return x
</code></pre>