Pruning and model size

username: lusamino
cooked: <p>Hi all,</p>
<p>This is my first question, so I will try to be not too clumsy. Besides, I am more or less a beginner with PyTorch, so I hope my questions are not that evident. And thanks for all your help beforehand!</p>
<p>Well, I just trained a simple model with 3 linear layers and everything went smooth and fine. Now, I want to replicate this model in order to allow some adversarial training of a few input samples at the same time, and for that I will be using the hack Goodfellow proposed <a href="https://github.com/tensorflow/tensorflow/issues/4897#issuecomment-290997283" rel="nofollow noopener">here</a>.</p>
<p>The problem is that the model may end up being quite big. But since many weights, all the off diagonal blocks, are going to be equal to zero, I thought: ok, let’s prune them, I don’t need them for anything. Then, the size of these repeated model will just scale with the number of replicas of the model, and not quadratically. However, I just put together a small example and this does not seem to work correctly:</p>
<pre><code class="lang-auto"># Original small model
enc_small = torch.nn.Sequential()
enc_small.add_module("dense_0", nn.Linear(2000, 4000))
torch.save(enc_small, 'orig_model.pth')

n_rep = 3

# Repeated model
enc_rep = torch.nn.Sequential()
enc_rep.add_module("dense_0", nn.Linear(2000 * n_rep, 4000 * n_rep))
torch.save(enc_rep, 'model_repeated.pth')

# Repeated model with pruning
enc_pruned = torch.nn.Sequential()
enc_pruned.add_module("dense_0", nn.Linear(2000 * n_rep, 4000 * n_rep))
# This mask is just zeroing out all elements out of the center blocks of 2000x4000
mask = np.zeros(encoder1.dense_0.weight.shape).astype(int)
for ir in range(n_rep + 1):
    mask[((ir - 1) * 4000):(4000*ir),
         ((ir - 1) * 2000):(2000*ir)] = 1

torch.nn.utils.prune.custom_from_mask(enc_pruned.dense_0, name='weight', mask = torch.tensor(mask))

torch.save(enc_pruned, 'pruned_model.pth')
</code></pre>
<p>And here is how the models look one saved:</p>
<pre><code class="lang-auto">25629256 -rw-r--r-- 1 users 1.4G Aug 16 09:47 pruned_model.pth
25629254 -rw-r--r-- 1 users 275M Aug 16 09:46 model_repeated.pth
25629253 -rw-r--r-- 1 users  31M Aug 16 09:46 orig_model.pth
</code></pre>
<p>I do believe I am miserably missing something, but I don’t know what… It might be the behaviour of <code>torch.nn.utils.prune.custom_from_mask</code>, that I am not understanding… Hence, I will really appreciate if you could offer me some help. Thank you so much!!!</p>
<p>P.S: I am in parallel experimenting with assigning the parameters as sparse matrices, but so far, to no avail…</p>

username: ptrblck
cooked: <p>I’m unfortunately not familiar enough with the pruning utilities, but <a class="mention" href="/u/michela">@Michela</a> might know, if the file size is expected.</p>
<p>PS: I would also generally recommend to store and load the <code>state_dict</code> instead of the complete model as described <a href="https://pytorch.org/docs/stable/notes/serialization.html">here</a>.</p>

username: lusamino
cooked: <p>Thank you so much <a class="mention" href="/u/ptrblck">@ptrblck</a>! Well, I believe <a class="mention" href="/u/michela">@Michela</a> will know <img src="https://discuss.pytorch.org/images/emoji/apple/slight_smile.png?v=9" title=":slight_smile:" class="emoji" alt=":slight_smile:"> Regarding what you mention, I will try then using <code>state_dict</code>. But it was already surprising to me that the size was changing that much, as on top of the <code>weight_mask</code>, I cannot think about any other overhead created on the pruned model.</p>

username: lusamino
cooked: <p>Ok, here I come with updated results. They are already better with <code>state_dict()</code>, but the pruned model still takes more memory. I just added:<br>
<code>torch.save(enc_pruned.state_dict(), 'pruned_model.pth')</code><br>
to all the models being saved.</p>
<p>And this is the output:</p>
<pre><code class="lang-auto">25629256 -rw-r--r-- 1 users 550M Aug 18 08:43 pruned_model.pth
25629254 -rw-r--r-- 1 users 275M Aug 18 08:43 model_repeated.pth
25629253 -rw-r--r-- 1 users  31M Aug 18 08:43 orig_model.pth
</code></pre>
<p>So, since the size is precisely twice the size of the original repeated, might it come from the mask? Is it then pruning correctly working?</p>
<p>Thanks again for all your help!!</p>

username: PabloVD
cooked: <p>Hi, I’m new with pruning and not really very familiar with this practice, but let me share my thoughts.</p>
<p>According to the <a href="https://pytorch.org/tutorials/intermediate/pruning_tutorial.html" rel="noopener nofollow ugc">pruning tutorial in pytorch</a>, applying a pruning routine like <code>torch.nn.utils.prune.custom_from_mask</code> only creates the masks, increasing the number of tensors stored in the model, and thus increasing its size. To really apply such masks, next one has to remove the pruned connections by running <code>prune.remove(module, "weight")</code> or similar.</p>
<p>In that case, saving with <code>torch.save</code> leads to stored models of the same size. However, following <a href="https://github.com/spellml/resnext50-panda/blob/master/notebooks/pruning.ipynb" rel="noopener nofollow ugc">this notebook</a> (which is explained in <a href="https://spell.ml/blog/model-pruning-in-pytorch-X9pXQRAAACIAcH9h" rel="noopener nofollow ugc">this tutorial</a>), if you compress the model, the pruned model becomes lighter than the original one.</p>
<p>Hope it helps!</p>