Model not training when implementing a custom loss

username: thomas
cooked: <p>I have a code of a custom loss that I implemented on tensorflow that I would like to pass on Pytorch for technical reasons. However I can’t seem to make it work and I don’t know why.<br>
The pytorch loss doesn’t seem to train the network</p>
<p>The aim is to compare raw linear output of a network (before softmax) with true probabilities.<br>
The wanted loss is similar to the one in <a href="https://www.youtube.com/watch?v=2vJtbAha3To" rel="nofollow noopener">this video</a> at 13:10</p>
<p>Any help would be appreciated <img src="https://discuss.pytorch.org/images/emoji/apple/slight_smile.png?v=9" title=":slight_smile:" class="emoji" alt=":slight_smile:"></p>
<p>Tensorflow code:</p>
<pre><code class="lang-auto">import numpy as np
from scipy.special import softmax

import tensorflow as tf
from tensorflow.keras import backend as K
from tensorflow.keras.optimizers import Adam
from tensorflow.keras.layers import Dense, Softmax
from tensorflow.keras import Input, Model

delta = tf.Variable([[1.]], trainable=False)

main_input = Input(shape=(10,))
output = Dense(4, activation='linear')(main_input)

def custom_loss(delta):
    def loss(y_true, y_pred):
        y_pred_softmax = Softmax()(y_pred)
        y_pred_softmax_clipped = K.clip(y_pred_softmax, 1e-8, 1 - 1e-8)
        log_likelihood = y_true * K.log(y_pred_softmax_clipped)
        return K.sum(-log_likelihood * delta)
    return loss

model = Model(inputs=[main_input], outputs=output)
model.compile(optimizer=Adam(lr=0.01), loss=custom_loss(delta))

print(model.predict(np.ones((1,10))))
print(softmax(model.predict(np.ones((1,10)))[0]))
delta.assign([[1.0]])
model.fit(np.ones((1000,10), dtype='float'),np.asarray(1000*[[0.7, 0.3, 0.0, 0.0]], dtype='float'))
print(model.predict(np.ones((1,10))))
print(softmax(model.predict(np.ones((1,10)))[0]))
</code></pre>
<p>Pytorch code:</p>
<pre><code class="lang-auto">import numpy as np

import torch
import torch.nn as nn
import torch.optim as optim
import torch.nn.functional as F


class Net(nn.Module):
    def __init__(self):
        super(Net, self).__init__()
        self.fc = nn.Linear(10, 4)

    def forward(self, x):
        x = self.fc(x)
        return x

def custom_loss(delta):
    def loss(y_pred, y_true):
        y_pred_softmax = nn.Softmax(dim=1)(y_pred)
        y_pred_softmax_clipped = torch.clamp(y_pred, 1e-8, 1 - 1e-8)
        log_likelihood = y_true * torch.log(y_pred_softmax_clipped)
        return torch.sum(-log_likelihood * delta)
    return loss


delta = 1

network = Net()
loss_function = custom_loss(delta)
optimizer = optim.SGD(network.parameters(), lr=0.01, momentum=0.9)
optimizer.zero_grad()


x = torch.ones((1,10))
print(network(x))
print(nn.Softmax(dim=1)(network(x)), '\n')

n=1000
x = torch.ones((n,10))
target = torch.from_numpy(np.array(n*[[0.8, 0.2, 0.0, 0.0]]))
output = network(x)

loss = loss_function(output, target)
loss.backward()
optimizer.step()
print(loss, '\n')

x = torch.ones((1,10))
print(network(x))
print(nn.Softmax(dim=1)(network(x)))
</code></pre>
<p>example of outputs, either only one value become extremely big :</p>
<pre><code class="lang-auto">tensor([[-0.0858,  0.1533, -0.1739,  2.0263]], grad_fn=&lt;AddmmBackward&gt;)
tensor([[0.0873, 0.1109, 0.0800, 0.7218]], grad_fn=&lt;SoftmaxBackward&gt;) 

tensor([[-8.5779e-02,  1.4366e+02, -1.7395e-01,  2.0263e+00]], grad_fn=&lt;AddmmBackward&gt;)
tensor([[0., 1., 0., 0.]], grad_fn=&lt;SoftmaxBackward&gt;)
</code></pre>
<p>Or the outputs doesn’t change at all:</p>
<pre><code class="lang-auto">tensor([[ 1.2120, -0.1411, -0.5820, -0.6478]], grad_fn=&lt;AddmmBackward&gt;)
tensor([[0.6327, 0.1635, 0.1052, 0.0985]], grad_fn=&lt;SoftmaxBackward&gt;) 

tensor([[ 1.2120, -0.1411, -0.5820, -0.6478]], grad_fn=&lt;AddmmBackward&gt;)
tensor([[0.6327, 0.1635, 0.1052, 0.0985]], grad_fn=&lt;SoftmaxBackward&gt;)
</code></pre>

username: KFrank
cooked: <p>Hi Thomas -</p>
<aside class="quote no-group" data-username="thomas" data-post="1" data-topic="91963" data-full="true">
<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/t/e36b37/48.png" class="avatar"> thomas:</div>
<blockquote>
<p>The pytorch loss doesn’t seem to train the network</p>
</blockquote>
</aside>
<p>The short answer is that your pytorch code only takes one optimizer<br>
step.</p>
<aside class="quote no-group">
<blockquote>
<p>Tensorflow code:</p>
<pre><code class="lang-auto">model.fit(np.ones((1000,10), dtype='float'),np.asarray(1000*[[0.7, 0.3, 0.0, 0.0]], dtype='float'))
</code></pre>
</blockquote>
</aside>
<p>With default arguments, <code>tf.keras.Model.fit()</code> performs one epoch<br>
of training with a batch size of 32.  You pass in 1000 samples, so you<br>
will train for about 30 optimizer steps.  This should be enough to train<br>
your very simple model, at least to some degree.</p>
<p>(Note, I have not looked at the rest of your code in any detail.)</p>
<aside class="quote no-group">
<blockquote>
<p>Pytorch code:</p>
<pre><code class="lang-auto">optimizer.step()
</code></pre>
</blockquote>
</aside>
<p>In contrast, <code>torch.optim.SGD.step()</code> takes only a single step of the<br>
optimizer.  In pytorch, you have to write your own (very simple) optimizer<br>
loop.</p>
<p>Good luck.</p>
<p>K. Frank</p>

username: thomas
cooked: <p>Oh thank you so much, I didn’t realize this mechanic about optimization before your comment, it made me understand many things about Pytorch</p>
<p>I tried to change my code, according to what you said in order to make the two code as similar as I possibly can, but I don’t know why it is still not working</p>
<pre><code class="lang-auto">import numpy as np

import torch
import torch.nn as nn
import torch.optim as optim
import torch.nn.functional as F


class Net(nn.Module):
    def __init__(self):
        super(Net, self).__init__()
        self.fc = nn.Linear(10, 4)

    def forward(self, x):
        x = self.fc(x)
        return x

    
def custom_loss(delta):
    def loss(y_pred, y_true):
        y_pred_softmax = nn.Softmax(dim=1)(y_pred)
        y_pred_softmax_clipped = torch.clamp(y_pred, 1e-8, 1 - 1e-8)
        log_likelihood = y_true * torch.log(y_pred_softmax_clipped)
        return torch.sum(-log_likelihood * delta)
    return loss


delta = 1
batch_size = 32
n_sample = 1000
network = Net()
loss_function = custom_loss(delta)
optimizer = optim.Adam(network.parameters(), lr=0.01)

x = torch.ones((1,10))
print(network(x))
print(nn.Softmax(dim=1)(network(x)), '\n')

target = [0.7, 0.3, 0.0, 0.0]

for i in range(int(n_sample/batch_size)):
    optimizer.zero_grad()
    
    inputs = torch.ones((batch_size,10))
    targets = torch.FloatTensor(batch_size*[target])
    outputs = network(inputs)
    
    loss = loss_function(outputs, targets)
    loss.backward()
    optimizer.step()

x = torch.ones((1,10), dtype=torch.float)
print(network(x))
print(nn.Softmax(dim=1)(network(x)))
</code></pre>
<p>I changed from SGD to Adam, just in case that would affect something (even if I doubt it would) as well as removing the momentum. I used the <a href="https://pytorch.org/tutorials/beginner/blitz/cifar10_tutorial.html#train-the-network" rel="nofollow noopener">example</a> of the pytorch website to be sure that I was not missing a step in the training procedure.<br>
However it still doesn’t work as intended <img src="https://discuss.pytorch.org/images/emoji/apple/confused.png?v=9" title=":confused:" class="emoji" alt=":confused:"></p>
<p>Two examples of outputs:</p>
<pre><code class="lang-auto">tensor([[-0.5445,  1.1960, -0.1856,  0.3652]], grad_fn=&lt;AddmmBackward&gt;)
tensor([[0.0942, 0.5370, 0.1349, 0.2340]], grad_fn=&lt;SoftmaxBackward&gt;) 

tensor([[-0.5445,  1.1960, -0.1856,  0.3652]], grad_fn=&lt;AddmmBackward&gt;)
tensor([[0.0942, 0.5370, 0.1349, 0.2340]], grad_fn=&lt;SoftmaxBackward&gt;)
</code></pre>
<pre><code class="lang-auto">tensor([[-0.1185,  0.5827,  0.7683,  0.4973]], grad_fn=&lt;AddmmBackward&gt;)
tensor([[0.1371, 0.2764, 0.3327, 0.2538]], grad_fn=&lt;SoftmaxBackward&gt;) 

tensor([[-0.1185,  1.8670,  0.7683,  0.4973]], grad_fn=&lt;AddmmBackward&gt;)
tensor([[0.0796, 0.5798, 0.1932, 0.1474]], grad_fn=&lt;SoftmaxBackward&gt;)
</code></pre>
<p>output with the TF code:</p>
<pre><code class="lang-auto">[[ 0.0284583   0.35027373  0.4858752  -0.61461353]]
[0.2229509  0.30758977 0.35225958 0.11719974]
Train on 1000 samples
1000/1000 [==============================] - 0s 302us/sample - loss: 25.1664
[[ 2.3004684  1.4706173 -1.8008366 -2.9013252]]
[0.68579024 0.29908288 0.01135056 0.00377643]
</code></pre>
<p>I also tried to increase the number of sample and modify the learning rate, but the results were still the same</p>
<p>Did you see other mistake in my code that could cause that ?</p>
<p>(again, thank you <img src="https://discuss.pytorch.org/images/emoji/apple/slight_smile.png?v=9" title=":slight_smile:" class="emoji" alt=":slight_smile:"> )</p>
<p>edit: I checked the forward and backward graph and all the function of the loss function are present except softmax, is it normal ?</p>
<pre><code class="lang-auto">from torchviz import make_dot, make_dot_from_trace
make_dot(loss, params=dict(network.named_parameters()))
</code></pre>

username: KFrank
cooked: <p>Hi Thomas!</p>
<aside class="quote no-group" data-username="thomas" data-post="3" data-topic="91963" data-full="true">
<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/t/e36b37/48.png" class="avatar"> thomas:</div>
<blockquote>
<p>I checked the forward and backward graph and all the function of the loss function are present except softmax, is it normal ?</p>
</blockquote>
</aside>
<p>This is the hint.</p>
<aside class="quote no-group">
<blockquote>
<pre><code class="lang-auto">def custom_loss(delta):
    def loss(y_pred, y_true):
        y_pred_softmax = nn.Softmax(dim=1)(y_pred)
        y_pred_softmax_clipped = torch.clamp(y_pred, 1e-8, 1 - 1e-8)
        log_likelihood = y_true * torch.log(y_pred_softmax_clipped)
        return torch.sum(-log_likelihood * delta)
    return loss
</code></pre>
</blockquote>
</aside>
<p>You pass <code>y_pred</code> (rather than <code>y_pred_softmax</code>) to <code>torch.clamp()</code>.<br>
This looks like a simple typo, and the effect is that the call to<br>
<code>nn.Softmax()</code> is bypassed.  (Note, you don’t have this typo in your<br>
original tensorflow code.)</p>
<p>Best.</p>
<p>K. Frank</p>

username: thomas
cooked: <p>Well, that was embarrassing… <img src="https://discuss.pytorch.org/images/emoji/apple/sweat_smile.png?v=9" title=":sweat_smile:" class="emoji" alt=":sweat_smile:"> <img src="https://discuss.pytorch.org/images/emoji/apple/no_mouth.png?v=9" title=":no_mouth:" class="emoji" alt=":no_mouth:"><br>
Thank you for your help! It’s working perfectly now <img src="https://discuss.pytorch.org/images/emoji/apple/slight_smile.png?v=9" title=":slight_smile:" class="emoji" alt=":slight_smile:"></p>
<p>Best,<br>
Thomas</p>