Pytorch pruning tutorial. Where in process does this happen?

jwillette: <p>I was looking at the Pytorch pruning tutorial (<a href="https://pytorch.org/tutorials/intermediate/pruning_tutorial.html#" rel="noopener nofollow ugc">https://pytorch.org/tutorials/intermediate/pruning_tutorial.html#</a>) and I am left confused about where in the process this pruning happens.</p>
<ol>
<li>Does the setup in the tutorial happen before training, and then it is trained with the pruning method somewhere in the loop</li>
<li>Do the strategies in the tutorial happen after training when we have the final network weights?</li>
</ol>

albanD: <p>Hi,</p>
<p>This tutorial shows how to use pruning for the weights that are used during training.<br>
It is not a pre or post processing technique.</p>
<p>In particular, the prunning happen when you access the weights where you registered the prunning, they will automatically get prunned.</p>

jwillette: <p>Thanks, so just to verify, I would follow the tutorial, and then make a normal training loop without any other modification and the model will automatically be pruned while training?</p>

albanD: <p>Yes this is the goal!</p>

Peter78: <p>Thanks,<br>
but I was wondering if the loss function is calculated using pruned tensors or unpruned ones.</p>
<p>Regards,</p>

albanD: <p>After the hooks are setup, everything is computed with the pruned Tensors.</p>

Peter78: <p>My question was related to the following discussion:</p>
<aside class="quote quote-modified" data-post="6" data-topic="93395">
  <div class="title">
    <div class="quote-controls"></div>
    <img loading="lazy" alt="" width="24" height="24" src="https://discuss.pytorch.org/user_avatar/discuss.pytorch.org/alband/48/215_2.png" class="avatar">
    <a href="https://discuss.pytorch.org/t/prunning-weights-not-optimizing-some-part-of-weights/93395/6">Prunning weights (not optimizing some part of weights)</a> <a class="badge-category__wrapper " href="/c/autograd/7"><span data-category-id="7" style="--category-badge-color: #AB9364; --category-badge-text-color: #FFFFFF;" data-drop-close="true" class="badge-category " title="A category of posts relating to the autograd engine itself."><span class="badge-category__name">autograd</span></span></a>
  </div>
  <blockquote>
    But, I don’t know how this can be implemented in pytorch. 

I think it depends a lot on the behavior you want. 
But just setting part of the Tensor to 0 before the forward and again after the optimizer’s step would work just fine if it’s what you want. Or you could also set some gradients to 0 before doing the optimizer step (be careful as weight decay/momentum might make parameters with 0 grad still change). 
But that depends how your pruning change with the iterations. 
I think you can do jus…
  </blockquote>
</aside>

<p>I would like to understand if the flow you proposed, that is to say:</p>
<pre><code class="lang-auto">for sample in dataset:
    model.prune_some_weights()
    out = model(sample)
    loss = crit(out)
    opt.zero_grad()
    loss.backward()
    opt.step()
    model.prune_some_weights()
</code></pre>
<p>it is still valid.</p>
<p>Regards</p>

albanD: <p>This version was for the case without hooks. In there indeed, you need to be a lot more careful to make sure the pruning is happening at the right time.<br>
The hooks remove the need for this fine balancing by pruning on the fly whenever you ask for the weights (in particular the reparametrization hooks).</p>