Optimizing a mask for weights instead of weights themselves

username: Soroush_Hashemifar
cooked: <p>Hi. I’m trying to optimize a mask for each weight of a simple DNN without updating the weights directly. My goal is to find a mask than I can multiply to network weights and achieve the maximum accuracy. This mask is the same size of the weights and denotes a coefficient of each weight parameter. The code I implemented is as follows:</p>
<pre><code class="lang-auto">learning_rate = 1
num_iterations = 3

weights = list(model.parameters())
masks = [torch.rand_like(layer.data, requires_grad=True) for layer in model.parameters()]

optimizer = torch.optim.Adam(indicators, lr=learning_rate)
criterion = torch.nn.CrossEntropyLoss()

for _ in range(num_iterations):
    for x, y in self.test_loader:
        optimizer.zero_grad()

        for idx, param in enumerate(model.parameters()):
            param = param * masks[idx]

        x = model(x)
        loss = criterion(x, y)

        loss.backward()

        print("mask", masks[-1].data)
        print("grad", masks[-1].grad)

        optimizer.step()

        with torch.no_grad():
            masks = [mask.clamp_(0.0, 1.0) for mask in masks]
</code></pre>
<p>I multiply the masks by the weights and then feed a data to the model. then backward the loss and since only masks are grad required, I expect it to calculate the gradients of masks. But this code gives “None” for all gradients of masks.<br>
I’ve spent lots of time to resolve it but haven’t been successful yet.</p>

username: KFrank
cooked: <p>Hi Soroush!</p>
<p>This is more of a pure python issue than one of pytorch in that you are<br>
setting a python name (a “variable”) to refer to a new object, rather than<br>
modifying the object to which that name previously referred.</p>
<p>The short story is that you are not actually modifying your model weights,<br>
so backpropagating through <code>model</code> doesn’t compute the gradients you<br>
expect.</p>
<aside class="quote no-group" data-username="Soroush_Hashemifar" data-post="1" data-topic="181834" data-full="true">
<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/soroush_hashemifar/48/61176_2.png" class="avatar"> Soroush_Hashemifar:</div>
<blockquote>
<p>I’m trying to optimize a mask for each weight of a simple DNN without updating the weights directly.</p>
</blockquote>
</aside>
<p>As as aside, why?  Why not optimize your model weights directly?</p>
<aside class="quote no-group">
<blockquote>
<pre><code class="lang-auto">        for idx, param in enumerate(model.parameters()):
</code></pre>
</blockquote>
</aside>
<p>At this point <code>param</code> is a name in your python script that refers to one of the<br>
parameters in your model.</p>
<aside class="quote no-group">
<blockquote>
<pre><code class="lang-auto">            param = param * masks[idx]
</code></pre>
</blockquote>
</aside>
<p>On the right-hand side of the equals sign a new tensor is created by multiplying<br>
together the model parameter referred to by <code>param</code> and <code>masks[idx]</code>.  The<br>
equals sign then sets the name <code>param</code> to refer to this <em>new</em> tensor.  After this<br>
line of code, <code>param</code> no longer refers to a parameter of your model.</p>
<p>Because <code>masks[idx]</code> has <code>requires_grad = True</code> (as presumably do your<br>
model parameters), the new tensor to which <code>param</code> now refers also has<br>
<code>requires_grad = True</code>, but this doesn’t matter, because you never use<br>
(the tensor referred to by) <code>param</code> again.</p>
<aside class="quote no-group">
<blockquote>
<pre><code class="lang-auto">        x = model(x)
        loss = criterion(x, y)
</code></pre>
</blockquote>
</aside>
<p><code>x</code> and <code>loss</code> depend on your model parameters, but they do <em>not</em> depend<br>
on <code>masks</code> because your model parameters were not updated to depend<br>
on <code>masks</code>.  That is, your model does not depend of the new tensor created<br>
by <code>param * masks[idx]</code>, so neither do <code>x</code> nor <code>loss</code>.</p>
<aside class="quote no-group">
<blockquote>
<pre><code class="lang-auto">        loss.backward()
</code></pre>
</blockquote>
</aside>
<p><code>loss</code> does not depend on <code>masks</code>, so <code>loss.backward()</code> does not compute<br>
gradients with respect to <code>masks</code>.</p>
<aside class="quote no-group">
<blockquote>
<p>I expect it to calculate the gradients of masks. But this code gives “None” for all gradients of masks.</p>
</blockquote>
</aside>
<p>As outlined above, gradients with respect to <code>masks</code> are not computed, so,<br>
indeed, <code>masks[-1].grad</code> remains <code>None</code>.</p>
<p>Best.</p>
<p>K. Frank</p>

username: Soroush_Hashemifar
cooked: <p>Thank you Frank for your thorough explanation. But to be honest I still don’t know how to fix it <img src="https://discuss.pytorch.org/images/emoji/apple/sweat_smile.png?v=12" title=":sweat_smile:" class="emoji" alt=":sweat_smile:" loading="lazy" width="20" height="20"><br>
To answer your first question</p>
<aside class="quote no-group" data-username="KFrank" data-post="2" data-topic="181834">
<div class="title">
<div class="quote-controls"></div>
<img loading="lazy" alt="" width="24" height="24" src="https://discuss.pytorch.org/letter_avatar_proxy/v4/letter/k/ecb155/48.png" class="avatar"> KFrank:</div>
<blockquote>
<p>Why not optimize your model weights directly?</p>
</blockquote>
</aside>
<p>This code is part of a bigger code that aims to prune the network parameters. So this mask tensor[s] is going to be used in order to prune some parameters later. And I forgot to mention that the model weights are pretrained.</p>
<p>Second, you’re correct that the mask is not directly used in the model according to the provided code snippet, so I can’t backward the gradient of loss w.r.t it. But, I’m following the notion of backpropagation. Please refine me if it’s not correct. In an ordinary case, we feed forward like this:</p>
<blockquote>
<p>A = ReLU(W . X + b)</p>
</blockquote>
<p>but in my case, I’m calculating it like this:</p>
<blockquote>
<p>A’ = ReLU((M * W) . X + b)</p>
</blockquote>
<p>so I think masks M is somehow contributing in the model, so I can calculate gradients w.r.t M. But I don’t know how exactly!</p>

username: KFrank
cooked: <p>Hi Soroush!</p>
<aside class="quote no-group" data-username="Soroush_Hashemifar" data-post="3" data-topic="181834" data-full="true">
<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/soroush_hashemifar/48/61176_2.png" class="avatar"> Soroush_Hashemifar:</div>
<blockquote>
<p>the mask is not directly used in the model according to the provided code snippet, so I can’t backward the gradient of loss w.r.t it.</p>
</blockquote>
</aside>
<p>Just to be clear <em>why</em> <code>masks</code> are not used in the model:</p>
<p>As mentioned above, this code:</p>
<pre data-code-wrap="python"><code class="lang-python">        for idx, param in enumerate(model.parameters()):
            param = param * masks[idx]
</code></pre>
<p>creates new tensors but does not modify the <code>model.parameters()</code> tensors.</p>
<p>In this context, <code>param = param * masks[idx]</code> is effectively equivalent to<br>
<code>new_tensor_that_will_be_ignored and_discarded = param * masks[idx]</code>.</p>
<p>Consider:</p>
<pre><code class="lang-plaintext">&gt;&gt;&gt; import torch
&gt;&gt;&gt; print (torch.__version__)
2.0.0
&gt;&gt;&gt;
&gt;&gt;&gt; _ = torch.manual_seed (2023)
&gt;&gt;&gt;
&gt;&gt;&gt; lin = torch.nn.Linear (2, 3)                     # a pytorch "model"
&gt;&gt;&gt;
&gt;&gt;&gt; lin.weight                                       # its first parameter
Parameter containing:
tensor([[-0.1004,  0.3112],
        [ 0.6338, -0.0288],
        [ 0.0585,  0.6938]], requires_grad=True)
&gt;&gt;&gt; lin.bias                                         # its second parameter
Parameter containing:
tensor([-0.1293, -0.3984, -0.4478], requires_grad=True)
&gt;&gt;&gt;
&gt;&gt;&gt; m1 = torch.randn (3, 2, requires_grad = True)    # first mask
&gt;&gt;&gt; m2 = torch.randn (3, requires_grad = True)       # second mask
&gt;&gt;&gt; masks = [m1, m2]                                 # packaged as a list
&gt;&gt;&gt;
&gt;&gt;&gt; x = torch.randn (2)                              # some input to the model
&gt;&gt;&gt;
&gt;&gt;&gt; lin (x)                                          # output before attempted modification
tensor([-0.4202, -0.0696, -0.9600], grad_fn=&lt;AddBackward0&gt;)
&gt;&gt;&gt;
&gt;&gt;&gt; for idx, param in enumerate(lin.parameters()):   # failed attempt at modification
...     param = param * masks[idx]                   # creates and discards new tensors
...
&gt;&gt;&gt; lin.weight                                       # model parameters unchanged
Parameter containing:
tensor([[-0.1004,  0.3112],
        [ 0.6338, -0.0288],
        [ 0.0585,  0.6938]], requires_grad=True)
&gt;&gt;&gt; lin.bias                                         # model parameters unchanged
Parameter containing:
tensor([-0.1293, -0.3984, -0.4478], requires_grad=True)
&gt;&gt;&gt;
&gt;&gt;&gt; lin (x)                                          # output unchanged after attempted modification
tensor([-0.4202, -0.0696, -0.9600], grad_fn=&lt;AddBackward0&gt;)
&gt;&gt;&gt;
&gt;&gt;&gt; loss = lin (x).sum()                             # some dummy loss
&gt;&gt;&gt; loss.backward()                                  # backpropagate
&gt;&gt;&gt;
&gt;&gt;&gt; m1.grad                                          # no grads for masks
&gt;&gt;&gt; m2.grad                                          # no grads for masks
&gt;&gt;&gt;
&gt;&gt;&gt; lin.weight.grad                                  # but model parameters do have grads
tensor([[ 0.4833, -0.7791],
        [ 0.4833, -0.7791],
        [ 0.4833, -0.7791]])
&gt;&gt;&gt; lin.bias.grad                                    # but model parameters do have grads
tensor([1., 1., 1.])
</code></pre>
<aside class="quote no-group">
<blockquote>
<p>I think masks M is somehow contributing in the model, so I can calculate gradients w.r.t M. But I don’t know how exactly!</p>
</blockquote>
</aside>
<p>Here’s how to implement this so that <code>masks</code> does contribute to <code>model</code> and<br>
that you can compute gradients with respect to <code>masks</code>.</p>
<p>Normally, pytorch will not let you modify model parameters inplace, but if<br>
you freeze the parameters (as you might be doing for pretrained weights)<br>
by setting their <code>requires_grad = False</code>, you can.  Then you can multiply<br>
by <code>masks</code> inplace, <code>param.mul_ (masks[idx])</code>, and things will work:</p>
<pre><code class="lang-plaintext">&gt;&gt;&gt; lin.weight.grad = None                           # clear model grads
&gt;&gt;&gt; lin.bias.grad = None                             # clear model grads
&gt;&gt;&gt;
&gt;&gt;&gt; lin.weight.requires_grad = False                 # freeze model parameters
&gt;&gt;&gt; lin.bias.requires_grad = False                   # by setting requires_grad = True
&gt;&gt;&gt;
&gt;&gt;&gt; lin.weight                                       # requires_grad is False (so doesn't display)
Parameter containing:
tensor([[-0.1004,  0.3112],
        [ 0.6338, -0.0288],
        [ 0.0585,  0.6938]])
&gt;&gt;&gt; lin.bias                                         # requires_grad is False (so doesn't display)
Parameter containing:
tensor([-0.1293, -0.3984, -0.4478])
&gt;&gt;&gt;
&gt;&gt;&gt; for idx, param in enumerate(lin.parameters()):   # successful parameter modification
...     _ = param.mul_ (masks[idx])                  # modify actual parameters inplace
...
&gt;&gt;&gt; lin.weight                                       # model parameters modified
Parameter containing:
tensor([[ 0.0377,  0.5826],
        [-0.1272,  0.0178],
        [-0.0097, -0.5214]], grad_fn=&lt;MulBackward0&gt;)
&gt;&gt;&gt; lin.bias                                         # and have grad_fn because of modification
Parameter containing:
tensor([-0.1900, -0.1723,  0.0791], grad_fn=&lt;MulBackward0&gt;)
&gt;&gt;&gt;
&gt;&gt;&gt; lin (x)                                          # changed output reflects changed parameters
tensor([-0.6257, -0.2476,  0.4806], grad_fn=&lt;AddBackward0&gt;)
&gt;&gt;&gt;
&gt;&gt;&gt; loss = lin (x).sum()                             # some dummy loss
&gt;&gt;&gt; loss.backward()                                  # backpropagate
&gt;&gt;&gt;
&gt;&gt;&gt; m1.grad                                          # now gradients flow back to masks
tensor([[-0.0485, -0.2425],
        [ 0.3063,  0.0224],
        [ 0.0283, -0.5405]])
&gt;&gt;&gt; m2.grad                                          # now gradients flow back to masks
tensor([-0.1293, -0.3984, -0.4478])
</code></pre>
<p>Best.</p>
<p>K. Frank</p>

username: Soroush_Hashemifar
cooked: <p>Thank you so much <a class="mention" href="/u/kfrank">@KFrank</a> for your descriptive example. Now I know what’s going on <img src="https://discuss.pytorch.org/images/emoji/apple/pray.png?v=12" title=":pray:" class="emoji" alt=":pray:" loading="lazy" width="20" height="20"></p>