Custom loss function not decreasing

username: maroov
cooked: <p>Sorry for the previous (deleted) post, I accidentally created the post before writing (I still don’t know how, but whatever).</p>
<ul>
<li>My model:</li>
</ul>
<pre><code class="lang-auto">class Module1(nn.Module):

  def __init__(self):
    super(Module1, self).__init__()
    self.conv1 = nn.Conv2d(3, 32, 5, 1, 2)
    self.conv2 = nn.Conv2d(32, 64, 5, 1, 2)
    self.conv3 = nn.Conv2d(64, 128, 5, 1, 2)
    self.conv4 = nn.Conv2d(128, 128, 5, 1, 2)
    self.pool = nn.MaxPool2d(2, 2)
    self.fc = nn.Linear(128*8*8, 10)

  def forward(self, x):
    x = self.pool(F.relu(self.conv2(F.relu(self.conv1(x)))))
    x = self.pool(F.relu(self.conv4(F.relu(self.conv3(x)))))
    x = x.view(x.size(0), -1)
    return x
</code></pre>
<ul>
<li>My loss:</li>
</ul>
<pre><code class="lang-auto">def infoNCE(zt, ztk):
  ind = np.diag_indices(b_size)
  aux = torch.exp(ztk @ torch.t(zt))[ind[0], ind[1]]
  aux = aux / torch.sum(torch.exp(ztk @ torch.t(zt)), axis=0)
  val = -torch.sum(torch.log(aux))
  return val
</code></pre>
<ul>
<li>My code for the training:</li>
</ul>
<pre><code class="lang-auto">mod1 = Module1().to(device)
optimizer = optim.Adam(mod1.parameters(), lr=lr)

for epoch in trange(epochs):
  for i, data in enumerate(trainloader):
    xt = data[0][0].to(device)
    xtk = data[0][1].to(device)
    yt = data[1].to(device)  # ignore this
    optimizer.zero_grad()

    zt = mod1(xt)
    ztk = mod1(xtk)
    loss = infoNCE(zt, ztk)
    loss.backward()
    optimizer.step()
</code></pre>
<ul>
<li>My loss over time:</li>
</ul>
<pre><code class="lang-auto">Epoch 50	Loss = 44.24047088623047
Epoch 100	Loss = 44.24409484863281
Epoch 150	Loss = 44.193824768066406
Epoch 200	Loss = 44.24012756347656
Epoch 250	Loss = 44.26813507080078
Epoch 300	Loss = 44.28200149536133
</code></pre>
<p><img src="https://discuss.pytorch.org/uploads/default/original/3X/7/b/7be47eba611d549b0f109d7407c37e4875bcd331.png" alt="image" data-base62-sha1="hG0hbBrRM4mELYz0dit5V81Prr3" width="391" height="248"></p>

username: charan_Vjy
cooked: <p>It is first important to localize the error. You could try one of two things.</p>
<ul>
<li>First try different learning rates, while simultaneously turning of all regularization.</li>
<li>Try using a standard loss function like the MSE (for regression) or the Cross Entropy (if classes are present). See if these loss fucntions decrease for a particular learning rate. If these losses do not decrease, it may indicate some underlying problem with the data or the way it was pre-processed.</li>
</ul>

username: braindotai
cooked: <p>I think there might be a bug, as you are not using <code>yt</code> anywhere after defining it.</p>

username: maroov
cooked: <p>That’s not a bug, sorry, I should have explained. My model uses  <code>xt</code> , which are patches of images, and  <code>xtk</code> , which are ‘future’ patches or predictions, that is, a different random patch of the same image but with some displacement. Variable  <code>yt</code>  represents the labels but it was meant to be there for an evaluation of the accuracy when testing, the labels are not used in this custom loss. Variable  <code>yt</code>  can be ignored in this post.</p>

username: braindotai
cooked: <p>Try lowering lr, check if that helps.</p>

username: maroov
cooked: <p>I can’t believe it was as simple as that, thank you</p>

username: braindotai
cooked: <p>Yeah, sometimes Adam don’t work with moderately high learning rate, so in that case I usually use sgd with low learning rate(0.001 or 0.01, these lrs are lower enough for sgd) to avoid nan with high momentum like 0.8 or 0.9 for fast convergence. You can try this as well.</p>

username: maroov
cooked: <p>Oh that’s sweet because I was about to say that I was getting NaN after several iterations, thanks again!</p>