Loss doesn't decrease in Pytorch CNN

First the major issues...

1. The main issue with this code is that you're using the wrong output shape and the wrong loss function for classification.

nn.BCELoss computes the binary cross entropy loss. This is applicable when you have one or more targets which are either 0 or 1 (hence the binary). In your case the target is a single integer between 0 and 9. Since there are only a small number of potential target values, the most common approach is to use categorical cross-entropy loss (nn.CrossEntropyLoss). The "theoretical" definition of cross entropy loss expects the network outputs and the targets to both be 10 dimensional vectors where the target is all zeros except in one location (one-hot encoded). However for computational stability and space efficiency reasons, pytorch's nn.CrossEntropyLoss directly takes the integer as a target. However, you still need to provide it with a 10 dimensional output vector from your network.

# pseudo code (ignoring batch dimension)
loss = nn.functional.cross_entropy_loss(<output 10d vector>, <integer target>)

To fix this issue in your code we need to have fc3 output a 10 dimensional feature, and we need the labels to be integers (not floats). Also, there's no need to use .sigmoid on fc3 since pytorch's cross-entropy loss function internally applies log-softmax before computing the final loss value.

2. As pointed out by Serget Dymchenko, you need to switch the network to eval mode during inference and train mode during train. This mainly affects dropout and batch_norm layers since they behave differently during training and inference.

3. A learning rate of 0.03 is probably a little too high. It works just fine with a learning rate of 0.001 and in a couple experiments I saw the training diverge at 0.03.


To accommodate these fixes a number of changes needed to be made. The minimal corrections to the code are shown below. I commented any lines which were changed with #### followed by a short description of the change.

import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F
import torch.optim as optim
from torch.utils.data import DataLoader
from torch.autograd import Variable
from keras.datasets import mnist

(x_train, y_train), (x_test, y_test) = mnist.load_data()


def resize(pics):
    pictures = []
    for image in pics:
        image = Image.fromarray(image).resize((dim, dim))
        image = np.array(image)
        pictures.append(image)
    return np.array(pictures)


dim = 60

x_train, x_test = resize(x_train), resize(x_test) # because my real problem is in 60x60

x_train = x_train.reshape(-1, 1, dim, dim).astype('float32') / 255
x_test = x_test.reshape(-1, 1, dim, dim).astype('float32') / 255
#### float32 -> int64
y_train, y_test = y_train.astype('int64'), y_test.astype('int64')

#### no reason to test for cuda before converting to numpy

#### I assume you were taking a subset for debugging? No reason to not use all the data
x_train = torch.from_numpy(x_train)
x_test = torch.from_numpy(x_test)
y_train = torch.from_numpy(y_train)
y_test = torch.from_numpy(y_test)


class ConvNet(nn.Module):

    def __init__(self):
        super().__init__()
        self.conv1 = nn.Conv2d(1, 32, 3)
        self.conv2 = nn.Conv2d(32, 64, 3)
        self.conv3 = nn.Conv2d(64, 128, 3)

        self.fc1 = nn.Linear(5*5*128, 1024)
        self.fc2 = nn.Linear(1024, 2048)
        #### 1 -> 10
        self.fc3 = nn.Linear(2048, 10)

    def forward(self, x):
        x = F.max_pool2d(F.relu(self.conv1(x)), (2, 2))
        x = F.max_pool2d(F.relu(self.conv2(x)), (2, 2))
        x = F.max_pool2d(F.relu(self.conv3(x)), (2, 2))

        x = x.view(x.size(0), -1)
        x = F.relu(self.fc1(x))
        x = F.relu(self.fc2(x))
        x = F.dropout(x, 0.5)
        #### removed sigmoid
        x = self.fc3(x)
        return x


net = ConvNet()

#### 0.03 -> 1e-3
optimizer = optim.Adam(net.parameters(), lr=1e-3)

#### BCELoss -> CrossEntropyLoss
loss_function = nn.CrossEntropyLoss()


class FaceTrain:

    def __init__(self):
        self.len = x_train.shape[0]
        self.x_train = x_train
        self.y_train = y_train

    def __getitem__(self, index):
        #### .unsqueeze(0) removed
        return x_train[index], y_train[index]

    def __len__(self):
        return self.len


class FaceTest:

    def __init__(self):
        self.len = x_test.shape[0]
        self.x_test = x_test
        self.y_test = y_test

    def __getitem__(self, index):
        #### .unsqueeze(0) removed
        return x_test[index], y_test[index]

    def __len__(self):
        return self.len


train = FaceTrain()
test = FaceTest()

train_loader = DataLoader(dataset=train, batch_size=64, shuffle=True)
test_loader = DataLoader(dataset=test, batch_size=64, shuffle=True)

epochs = 10
steps = 0
train_losses, test_losses = [], []
for e in range(epochs):
    running_loss = 0
    #### put net in train mode
    net.train()
    for idx, (images, labels) in enumerate(train_loader):
        optimizer.zero_grad()
        log_ps = net(images)
        loss = loss_function(log_ps, labels)
        loss.backward()
        optimizer.step()
        running_loss += loss.item()
    else:
        test_loss = 0
        accuracy = 0

        #### put net in eval mode
        net.eval()
        with torch.no_grad():
            for images, labels in test_loader:
                log_ps = net(images)
                test_loss += loss_function(log_ps, labels)
                #### removed torch.exp() since exponential is monotone, taking it doesn't change the order of outputs. Similarly with torch.softmax()
                top_p, top_class = log_ps.topk(1, dim=1)
                #### convert to float/long using proper methods. what you have won't work for cuda tensors.
                equals = top_class.long() == labels.long().view(*top_class.shape)
                accuracy += torch.mean(equals.float())
        train_losses.append(running_loss/len(train_loader))
        test_losses.append(test_loss/len(test_loader))
        print("[Epoch: {}/{}] ".format(e+1, epochs),
              "[Training Loss: {:.3f}] ".format(running_loss/len(train_loader)),
              "[Test Loss: {:.3f}] ".format(test_loss/len(test_loader)),
              "[Test Accuracy: {:.3f}]".format(accuracy/len(test_loader)))

Results of training are now...

[Epoch: 1/10]  [Training Loss: 0.139]  [Test Loss: 0.046]  [Test Accuracy: 0.986]
[Epoch: 2/10]  [Training Loss: 0.046]  [Test Loss: 0.042]  [Test Accuracy: 0.987]
[Epoch: 3/10]  [Training Loss: 0.031]  [Test Loss: 0.040]  [Test Accuracy: 0.988]
[Epoch: 4/10]  [Training Loss: 0.022]  [Test Loss: 0.029]  [Test Accuracy: 0.990]
[Epoch: 5/10]  [Training Loss: 0.017]  [Test Loss: 0.066]  [Test Accuracy: 0.987]
[Epoch: 6/10]  [Training Loss: 0.015]  [Test Loss: 0.056]  [Test Accuracy: 0.985]
[Epoch: 7/10]  [Training Loss: 0.018]  [Test Loss: 0.039]  [Test Accuracy: 0.991]
[Epoch: 8/10]  [Training Loss: 0.012]  [Test Loss: 0.057]  [Test Accuracy: 0.988]
[Epoch: 9/10]  [Training Loss: 0.012]  [Test Loss: 0.041]  [Test Accuracy: 0.991]
[Epoch: 10/10]  [Training Loss: 0.007]  [Test Loss: 0.048]  [Test Accuracy: 0.992]

Some other issues that will improve your performance and code.

4. You're never moving the model to the GPU. This means you won't be getting GPU acceleration.

5. torchvision is designed with all the standard transforms and datasets and is built to be used with PyTorch. I recommend using it. This also removes the dependency on keras in your code.

6. Normalize your data by subtracting the mean and dividing by the standard deviation to improve performance of your network. With torchvision you can use transforms.Normalize. This won't make a big difference in MNIST because its already too easy. But in more difficult problems it turns out to be important.


Further improved code is show below (much faster on GPU).

import torch
import torch.nn as nn
import torch.nn.functional as F
import torch.optim as optim
from torch.utils.data import DataLoader
from torchvision.datasets import MNIST
from torchvision import transforms

dim = 60

class ConvNet(nn.Module):
    def __init__(self):
        super().__init__()
        self.conv1 = nn.Conv2d(1, 32, 3)
        self.conv2 = nn.Conv2d(32, 64, 3)
        self.conv3 = nn.Conv2d(64, 128, 3)

        self.fc1 = nn.Linear(5 * 5 * 128, 1024)
        self.fc2 = nn.Linear(1024, 2048)
        self.fc3 = nn.Linear(2048, 10)

    def forward(self, x):
        x = F.max_pool2d(F.relu(self.conv1(x)), (2, 2))
        x = F.max_pool2d(F.relu(self.conv2(x)), (2, 2))
        x = F.max_pool2d(F.relu(self.conv3(x)), (2, 2))

        x = x.view(x.size(0), -1)
        x = F.relu(self.fc1(x))
        x = F.relu(self.fc2(x))
        x = F.dropout(x, 0.5)
        x = self.fc3(x)
        return x


net = ConvNet()
if torch.cuda.is_available():
    net.cuda()

optimizer = optim.Adam(net.parameters(), lr=1e-3)

loss_function = nn.CrossEntropyLoss()

train_dataset = MNIST('./data', train=True, download=True,
                      transform=transforms.Compose([
                          transforms.Resize((dim, dim)),
                          transforms.ToTensor(),
                          transforms.Normalize((0.1307,), (0.3081,))
                      ]))
test_dataset = MNIST('./data', train=False, download=True,
                     transform=transforms.Compose([
                         transforms.Resize((dim, dim)),
                         transforms.ToTensor(),
                         transforms.Normalize((0.1307,), (0.3081,))
                     ]))

train_loader = DataLoader(dataset=train_dataset, batch_size=64, shuffle=True, num_workers=8)
test_loader = DataLoader(dataset=test_dataset, batch_size=64, shuffle=False, num_workers=8)

epochs = 10
steps = 0
train_losses, test_losses = [], []
for e in range(epochs):
    running_loss = 0
    net.train()
    for images, labels in train_loader:
        if torch.cuda.is_available():
            images, labels = images.cuda(), labels.cuda()
        optimizer.zero_grad()
        log_ps = net(images)
        loss = loss_function(log_ps, labels)
        loss.backward()
        optimizer.step()
        running_loss += loss.item()
    else:
        test_loss = 0
        accuracy = 0

        net.eval()
        with torch.no_grad():
            for images, labels in test_loader:
                if torch.cuda.is_available():
                    images, labels = images.cuda(), labels.cuda()
                log_ps = net(images)
                test_loss += loss_function(log_ps, labels)
                top_p, top_class = log_ps.topk(1, dim=1)
                equals = top_class.flatten().long() == labels
                accuracy += torch.mean(equals.float()).item()
        train_losses.append(running_loss/len(train_loader))
        test_losses.append(test_loss/len(test_loader))
        print("[Epoch: {}/{}] ".format(e+1, epochs),
              "[Training Loss: {:.3f}] ".format(running_loss/len(train_loader)),
              "[Test Loss: {:.3f}] ".format(test_loss/len(test_loader)),
              "[Test Accuracy: {:.3f}]".format(accuracy/len(test_loader)))

Updated results of training...

[Epoch: 1/10]  [Training Loss: 0.125]  [Test Loss: 0.045]  [Test Accuracy: 0.987]
[Epoch: 2/10]  [Training Loss: 0.043]  [Test Loss: 0.031]  [Test Accuracy: 0.991]
[Epoch: 3/10]  [Training Loss: 0.030]  [Test Loss: 0.030]  [Test Accuracy: 0.991]
[Epoch: 4/10]  [Training Loss: 0.024]  [Test Loss: 0.046]  [Test Accuracy: 0.990]
[Epoch: 5/10]  [Training Loss: 0.020]  [Test Loss: 0.032]  [Test Accuracy: 0.992]
[Epoch: 6/10]  [Training Loss: 0.017]  [Test Loss: 0.046]  [Test Accuracy: 0.991]
[Epoch: 7/10]  [Training Loss: 0.015]  [Test Loss: 0.034]  [Test Accuracy: 0.992]
[Epoch: 8/10]  [Training Loss: 0.011]  [Test Loss: 0.048]  [Test Accuracy: 0.992]
[Epoch: 9/10]  [Training Loss: 0.012]  [Test Loss: 0.037]  [Test Accuracy: 0.991]
[Epoch: 10/10]  [Training Loss: 0.013]  [Test Loss: 0.038]  [Test Accuracy: 0.992]

One thing I noticed that you test the model in train mode. You need to call net.eval() to disable dropouts (and then net.train() again to put it back in the train mode).

Maybe there are other issues. Is training loss going down? Have you tried to overfit on a single example?