Custom loss function with trainable parameters

username: mark_eu
cooked: <p>Hi everyone,<br>
I need help.<br>
I am trying to create a custom loss function with two trainable parameters.</p>
<pre><code class="lang-auto">class MyCustomLoss(nn.Module):
    def __init__(self, my_parameter1, my_parameter2):
       super(MyCustomLoss, self).__init__()
	   self.A = my_parameter1
	   self.B = my_parameter2
    def forward(self, inputs, targets):
	   y_hat_softmax = F.softmax(inputs, dim=1)
	   t = torch.argmax(y_hat_softmax, dim=1)
	   
	   ...some custom code1...
	   
	   MyLoss = self.A*(some custom code2) + self.B*(some custom code3)
	   
	   return MyLoss
</code></pre>
<p>For the first prototyping of the idea, I used brute-force training of values A and B between 0 and 1<br>
with some small steps (0.05). I have proof that the idea works, but I would like to get the best<br>
values for A and B - that the network learns these weights.</p>
<p>Do you have any suggestions with the code snippet? Also, <strong>I would like to ensure that A and B are between 0 and 1</strong>, or at least, positive values.</p>
<p>Thank you</p>

username: Dazitu616
cooked: <p>To make A and B positive, an easy way is to apply ReLU to them before multiplying with the loss, i.e. <code>MyLoss = torch.relu(self.A)*(some custom code2) + torch.relu(self.B)*(some custom code3)</code>. Another option is to apply <code>torch.exp</code> to A and B, and this is a common trick people use in training VAE (to make the predicted variance positive)</p>
<p>To make it between 0, 1, similarly, you can apply a sigmoid on A and B</p>

username: mark_eu
cooked: <p><a class="mention" href="/u/dazitu616">@Dazitu616</a> thank you! It works with that approach.<br>
I would like to get the final value of the A and B at the end. How is it possible?</p>
<p>Thank you</p>

username: Dazitu616
cooked: <p>Just print it out? Or load the saved state_dict of the model and find their values?</p>

username: mark_eu
cooked: <p>Thank you, <a class="mention" href="/u/dazitu616">@Dazitu616</a>, for your suggestion.<br>
I added A and B in the optimizer in this way:</p>
<pre><code class="lang-auto">optimizer = torch.optim.Adam(list(model.parameters())+list(criterion.parameters()), lr=learning_rate)
</code></pre>
<p>It seems to me that printing state_dict is the easiest way. However, the following code does not list A and B:</p>
<pre><code class="lang-auto"># Print model's state_dict
print("Model's state_dict:")
for param_tensor in model.state_dict():
    print(param_tensor, "\t", model.state_dict()[param_tensor].size())

print()

# Print optimizer's state_dict
print("Optimizer's state_dict:")
for var_name in optimizer.state_dict():
    print(var_name, "\t", optimizer.state_dict()[var_name])
</code></pre>
<p>Do you have a suggestion on how to approach these values from the state_dict()?</p>
<p>Thank you</p>

username: Dazitu616
cooked: <p>Oh, sure. I think you need to manually include these two weights to the model by <a href="https://pytorch.org/docs/stable/generated/torch.nn.Module.html#torch.nn.Module.register_parameter" rel="noopener nofollow ugc">register_params</a>. For example:</p>
<pre><code class="lang-auto">class Model(nn.Module):

    def __init__(self):
        A = torch.tensor(0.5)
        self.register_parameter('A', A)
</code></pre>
<p>Then you can call it via <code>self.A</code>. And this <code>A</code> will be in the <code>model.state_dict()</code></p>

username: mark_eu
cooked: <p>Thank you, <a class="mention" href="/u/dazitu616">@Dazitu616</a>.<br>
I am a little bit confused. I have registered these parameters inside the custom loss class (as in the first post of my question). So why is there a need to repeat them inside the module class?<br>
I used this in my custom loss class:</p>
<pre><code class="lang-auto">self.A = torch.nn.Parameter(torch.tensor(0.0, requires_grad=True))
</code></pre>
<p>Thank you</p>

username: Dazitu616
cooked: <p>IC, sorry for the confusion. Since you have registered them inside the <strong>custom loss class</strong>, I think you can print the state_dict of the loss to see <code>A</code>, i.e. <code>criterion.state_dict()['A']</code></p>

username: mark_eu
cooked: <p>Thank you, <a class="mention" href="/u/dazitu616">@Dazitu616</a>.<br>
Yes, that is the correct syntax - I get the values. However, I get just the initial values (0.0). Do you have any idea where the catch is?</p>
<p>Thank you</p>

username: Dazitu616
cooked: <p>This is just my guess: since you are minimizing the loss using an optimizer, so if <code>A</code> goes to 0, then the loss will be 0, which is exactly the smallest possible value. So I guess it’s not a good idea of doing so (optimizing the loss scale). Instead, I think you can have some constraints, e.g. <code>A + B == 1</code>, so that the optimizer cannot find the trivial solution by setting <code>A</code> and <code>B</code> both to 0. Instead, it should find a balance which benefits the model. Just my random thought though.</p>

username: mark_eu
cooked: <p>Thank you, <a class="mention" href="/u/dazitu616">@Dazitu616</a>. Your suggestion was the correct one. The initial value can’t be 0.0 but some other value, 1.0, for instance.<br>
I have noticed that training gives a negative value for the A or B parameter for some initial values. So I used a code to have the loss always positive:</p>
<pre><code class="lang-auto">MyLoss = torch.relu(self.A)*(some custom code2) + torch.relu(self.B)*(some custom code3)
</code></pre>
<p>But is there a way to ensure that a raw value of A is just positive during the training?</p>

username: Dazitu616
cooked: <p>I’m afraid not. There is no constraints to make them positive right? Maybe you can try with different things <code>torch.relu</code>, <code>torch.abs</code>. But that’s another question.</p>

username: mark_eu
cooked: <p>Thank you, <a class="mention" href="/u/dazitu616">@Dazitu616</a>, for your suggestions.</p>