How to add trainable attributes outside the nn.Module entities?

Aryan_Asadian: <p>I have reviewed the forums and docs about adding a trainable parameter to the model, by defining <code>trainable_attr = nn.Parameter(None, requires_grad = True)</code> and after that register it to the target model <code>model.register_parameter(name, params)</code>. But How can I train a parameter outside of the model, for instance, I have a custom loss which is the addition of two losses, e.g., MSE and Cross_Entropy:<br>
<code>total_loss = mse_loss *gamma + cross_entroppy_loss *(1-gamma)</code><br>
How can I define and train gamma?</p>

Caruso: <p>Hi,</p>
<p>you can create a separate <code>nn.Module</code> which contains your loss functions and your learnable parameter, just like this:</p>
<pre><code class="lang-auto">class CostumLoss(nn.Module):
    def __init__(self):
        super().__init__()
        self.mse_fn = nn.MSELoss()
        self.ce_fn = nn.CrossEntropyLoss()
        self.gamma = nn.Parameter(torch.tensor([.5]))

    def forward(self, out, mse_target, ce_target):
        mse_loss = self.mse_fn(out, mse_target)
        ce_loss = self.ce_fn(out, ce_target)
        loss = mse_loss * self.gamma + ce_loss * (1 - self.gamma)
        return loss
</code></pre>
<p>For the optimization you can ether use a separate optimizer for your loss function or chain it with the model parameters:</p>
<pre><code class="lang-auto">import itertools

loss_fn = CostumLoss()
optim = optim.Adam(itertools.chain(model.parameters(), loss_fn.parameters()), lr=1e-3)
</code></pre>
<p>I think you could also declare <code>gamma</code> as a normal variable of type <code>nn.Parameter</code> outside of any Module, but i don’t know how you would ‘chain’ it with the other parameters, like i did above.</p>

Aryan_Asadian: <p>Hi,<br>
Thank you. It’s working without any error, but the model can not be trained, i.e., the loss doesn’t decrease.   In your code snippet, how do you involve the input x? I have this same question about your second solution as well, i.e, how can we make the chain between the input and our trainable gamma in this solution?</p>

Caruso: <p>Hi,</p>
<p>sorry I did a mistake, but its now corrected…so <code>x</code> was <code>out</code>. I don’t even know if this loss function would work tbh, because <code>gamma</code> is not required to be in range <code>[0, 1]</code> and could just shoot toward <code>+/- inf</code>.</p>
<p>For the second solution I don’t know how to correctly implement it myself…so you would write something like:</p>
<pre><code class="lang-auto">gamma = nn.Parameter(torch.tensor([.5]))
optim = optim.Adam(itertools.chain(model.parameters(), gamma), lr=1e-3)
</code></pre>
<p>but if you run this, it will give you an error, because iterating over gamma returns you a non-leaf tensor, see:</p>
<pre><code class="lang-auto">for g in gamma:
    print(g.is_leaf)
</code></pre>
<p>Hopefully someone can point out my error <img src="https://discuss.pytorch.org/images/emoji/apple/smiley.png?v=9" title=":smiley:" class="emoji" alt=":smiley:"></p>

drevicko: <p>I think you need something like this:</p>
<pre><code>itertools.chain(model.parameters(), (gamma,))
</code></pre>
<p>Otherwise you end up with gamma’s values (which aren’t Parameters).</p>