Optimize partial parameters in nn.Parameter

Dian-Shan: <p>Hello, I am trying to optimize the part of parameters in “one” nn.Parameter variable.</p>
<p>e.g. a simple CNN classifier</p>
<pre><code class="lang-auto">class Model(nn.Module):
    def __init__(self, device):
        super().__init__()
        self.resnet = ...
        
        feature_dim, n_cls = 64, 10
        self.classifier = torch.nn.Parameter(torch.Tensor(n_cls, feature_dim)).to(device)

    def forward(self, input):
        out = self.resnet(input)
        return out.matmul(self.classifier.t())

my_model = Model()
</code></pre>
<p>In some cases, I only want to optimize the certain rows in <code>self.classifier</code>.<br>
e.g. <code>[self.classifier[0], self.classifier[2], self.classifier[4], self.classifier[6], self.classifier[8]]</code></p>
<p>I figured out two possible ways to implement an optimizer for this situation, but both of them caused error.<br>
1.</p>
<pre><code class="lang-auto">params = []
for i, para in enumerate(my_model.classifier):
    if i % 2 == 0:
        params.append(para)

optimizer = Adam(params, ...)
</code></pre>
<p>error msg: TypeError: optimizer can only optimize Tensors, but one of the params is Module.parameters</p>
<ol start="2">
<li>another similar way but different parameter group for the optimizer</li>
</ol>
<pre><code class="lang-auto">params = []
for i, para in enumerate(my_model.classifier):
    if i % 2 == 0:
        params.append({"params": para})

optimizer = Adam(params, ...)
</code></pre>
<p>error msg: ValueError: can’t optimize a non-leaf Tensor</p>
<p>Does anyone have any ideas to optimize some entry in a nn.Parameter?</p>

ptrblck: <p>You could define separate parts of the <code>self.classifier</code> parameter and only pass the parts, which should be optimized, to the optimizer. In the <code>forward</code> method you would then have to recreate the “full” parameter using <code>torch.cat</code> and/or <code>torch.stack</code> and apply it in the <code>matmul</code>.<br>
Another approach would be to zero out the gradients of specific parts of the parameter, e.g. via <code>register_hook</code>, but this approach might be a bit tricky especially if your optimizer uses internal running stats.</p>

Dian-Shan: <p>Many thanks for your informative reply!<br>
I first thought about approach 2 (zero out the gradient). However, I am not familiar with <code>register_hook</code>, so implement this approach <strong>before</strong> <code>optimizer.step()</code>, I manually zero out the gradients. In this case, if I init the optimizer with passing <code>weight_decay</code> argument (not zero), the optimizer still makes the parameter updated. (ref <a href="https://github.com/pytorch/pytorch/blob/bb34fd6191e0f2a31edc38e00d93a24d9da8c031/torch/optim/_functional.py#L160" rel="noopener nofollow ugc">here</a>)</p>
<p>Therefore, I think approach 1 (separate parts of parameter) might be better in practice. <strong>Explicitly</strong> tell the optimizer which parameters should be updated. (I think set <code>requires_grad</code> is also good, checked <a href="https://github.com/pytorch/pytorch/blob/master/torch/optim/sgd.py#L100" rel="noopener nofollow ugc">here</a>)</p>
<p>Note: I am not sure whether <code>register_hook</code> can deal with the above mentioned update by <code>weight_decay</code>.</p>

ptrblck: <p>No, there wouldn’t be a difference in zeroing out the gradients by accessing the parameter directly or by using the hooks. For the mentioned shortcomings (or maybe unexpected behavior) I would also prefer the first approach.</p>

Dian-Shan: <p>Yes!<br>
I have implemented the first approach and it works well.<br>
Many thanks for your kindly help!</p>