Why is this nerual network performing poorly on MNIST?
Hi I am building a nerual network in pytorch to classify MNIST and for the life of me, I can't seem to find out why this network won't perform above 7% accuracy. Any guidance would be nice.
import torch
import torch.nn as nn
import torch.optim as optim
import torch.nn.functional as F
from keras.datasets import mnist
from keras.utils.np_utils import to_categorical
import numpy as np
from sklearn.metrics import precision_score, recall_score, f1_score, accuracy_score, confusion_matrix

(X_train, Y_train), (X_test, Y_test) = mnist.load_data()

X_train = X_train.astype(&quot;float32&quot;)/255
X_test = X_test.astype(&quot;float32&quot;)/255


X_train = X_train.reshape(X_train.shape[0],(X_train.shape[1] * X_train.shape[2]));
X_test = X_test.reshape(X_test.shape[0],(X_test.shape[1] * X_test.shape[2]));

class Net(torch.nn.Module):
  def __init__(self):
    super(Net, self).__init__()
    self.lin_1 = nn.Linear(784, 128)
    self.lin_2 = nn.Linear(128, 64)
    self.lin_3 = nn.Linear(64, 10)

  def forward(self,x) :
    x = self.lin_1(x)
    x = torch.relu(x)
    x = self.lin_2(x)
    x = torch.relu(x)
    x = self.lin_3(x)
    x = torch.softmax(x, dim=0)
    return x

net = Net();
loss = torch.nn.CrossEntropyLoss();
optimizer = torch.optim.SGD(net.parameters(),lr = 0.01);

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)

device = torch.device(&quot;cuda:0&quot; if torch.cuda.is_available() else &quot;cpu&quot;);
X_train.to(device);
X_test.to(device);
y_train.to(device);
y_test.to(device);
net.to(device);
loss.to(device);

y_train = y_train.type(torch.long)
y_test = y_test.type(torch.long)

net.train()
for epoch in range(10):
  #pred = torch.max(net(X_train),1);
  pred = net(X_train.to(device));
  
  train_loss = loss(pred,y_train.to(device));
  optimizer.zero_grad()
  train_loss.backward()
  optimizer.step()

net.eval()

pred = torch.max(net(X_test.to(device)),1)[1];
print('The accuracy for pytorch is ' , accuracy_score(y_test.cpu().numpy(),pred.cpu().numpy()));

I feel like I have to transform the data somehow. That's why I am dividing the the training and testing data by 255, and the network excepts a float for the input and a long for the output.
Here is a numpy version I made without pytorch
from keras.datasets import mnist
from keras.utils.np_utils import to_categorical
(X_train, Y_train), (X_test, Y_test) = mnist.load_data()

X_train = X_train.astype(&quot;float32&quot;)/255
X_test = X_test.astype(&quot;float32&quot;)/255


X_train = X_train.reshape(X_train.shape[0],(X_train.shape[1] * X_train.shape[2]));
X_test = X_test.reshape(X_test.shape[0],(X_test.shape[1] * X_test.shape[2]));

Y_train = to_categorical(Y_train);
Y_test = to_categorical(Y_test)
import numpy as np

print(Y_test.shape)

class DNN():
  def __init__(self, sizes, epochs=10, lr = 0.01):
    self.sizes = sizes
    self.epochs = epochs
    self.lr = lr
    self.params = self.initialization();

  def ReLu(self, x, derivative=False):
    if derivative:
      return 1. * (x &gt; 0)
    else:
      return x * (x &gt; 0)

  def softmax(self, x, derivative=False):
        # Numerically stable with large exponentials
        exps = np.exp(x - x.max())
        if derivative:
            return exps / np.sum(exps, axis=0) * (1 - exps / np.sum(exps, axis=0))
        return exps / np.sum(exps, axis=0)

  def initialization(self):
        # number of nodes in each layer
        input_layer=self.sizes[0]
        hidden_1=self.sizes[1]
        hidden_2=self.sizes[2]
        output_layer=self.sizes[3]

        params = {
            &quot;W1&quot;:np.random.randn(hidden_1, input_layer) * np.sqrt(1. / hidden_1),
            &quot;W2&quot;:np.random.randn(hidden_2, hidden_1) * np.sqrt(1. / hidden_2),
            &quot;W3&quot;:np.random.randn(output_layer, hidden_2) * np.sqrt(1. / output_layer)
        }

        return params
  def forward (self,X_train):
    
    self.params[&quot;X0&quot;] = X_train;
    
    self.params[&quot;Z1&quot;] = np.dot(self.params[&quot;W1&quot;], self.params[&quot;X0&quot;])
    self.params['X1'] = self.ReLu(self.params[&quot;Z1&quot;])

    self.params['Z2'] = np.dot(self.params[&quot;W2&quot;], self.params[&quot;X1&quot;])
    self.params[&quot;X2&quot;] = self.ReLu(self.params[&quot;Z2&quot;])

    self.params[&quot;Z3&quot;] = np.dot(self.params[&quot;W3&quot;], self.params[&quot;X2&quot;])
    self.params[&quot;X3&quot;] = self.softmax(self.params[&quot;Z3&quot;])

    return self.params[&quot;X3&quot;]
  
  def backpropagation (self, Y_train, output):

    update = {};

    error = 2 * (output - Y_train) / output.shape[0] * self.softmax(self.params[&quot;Z3&quot;], derivative=True)
    update[&quot;W3&quot;] = np.outer(error, self.params[&quot;X2&quot;])

    error = np.dot(self.params[&quot;W3&quot;].T, error) * self.ReLu(self.params[&quot;Z2&quot;], derivative=True)
    update[&quot;W2&quot;] = np.outer(error, self.params[&quot;X1&quot;])

    error = np.dot(self.params[&quot;W2&quot;].T, error) * self.ReLu(self.params[&quot;Z1&quot;], derivative=True)
    update[&quot;W1&quot;] = np.outer(error, self.params[&quot;X0&quot;])

    return update

  def updateParams (self,update):
    for key, value in update.items():
      #print(key)
      self.params[key] -= self.lr * value

  def test_accuracy(self, X_test, Y_train):
    predictions = []
    for i in range(len(X_test)):
      output = self.forward(X_test[i])
      pred = np.argmax(output)
      predictions.append(pred == np.argmax(Y_train[i]))
    
    
    return np.mean(predictions)


  def train(self, X_train, Y_train):
        for epoch in range(self.epochs):
            print(&quot;epoch &quot;, epoch)
            for i in range(len(X_train)):
                output = self.forward(X_train[i])
                update = self.backpropagation(Y_train[i], output)
                self.updateParams(update)

dnn = DNN(sizes=[784, 200, 50, 10],epochs=10)
dnn.train(X_train, Y_train)

print(&quot;The accuracy of the numpy network on the test dataset is &quot;, dnn.test_accuracy(X_test,Y_test))



Well, I can immediately tell there are a couple of problems with the code you provided:

Please check the documentation for PyTorch's cross entropy loss function. If you read it, you'll notice that torch.nn.CrossEntropyLoss performs the softmax function internally. This means that you shouldn't really be using another torch.softmax as the output activation if you're using nn.CrossEntropyLoss. If for some reason you want to use softmax at the output layer, you should consider using nn.NLLLoss instead. If you look at the image I posted below, simply removing x = torch.softmax(x, dim=0) causes the loss to fall, whereas using it causes the loss to be the same (hence, bad).

You are training with way too few epochs. I tried running your code with 3,000 epochs rather than 10, and the end performance is 0.9028 rather than the original 0.1038. You can also see that the loss value drops much more compared to the original implementation (second picture).




Edit
After taking a look at your NumPy code the problem became clearer. My second point still holds in essence: you're not training your model nearly enough. I somewhat incorrectly used the term &quot;epoch&quot; above, but what I really meant is &quot;steps.&quot;
If you look at your NumPy code, you have two for loops: the outer one is the number of epochs and the inner one loops through the training data. You're apparently using single-batch training for ten epochs. This implies that you're updating your model's parameters for a total of 600,000 times (60,000 training samples * 10 epochs) for the entire process. For your PyTorch code you're feeding the entire training data in one batch and training for ten epochs. This means you're updating your parameters only ten times.
If you modify your PyTorch code to be:
for epoch in range(10):
    net.train()

    for idx, _ in enumerate(X_train):
        prediction = net(X_train[idx].to(device))
        train_loss = loss(prediction.unsqueeze(0), y_train[idx].unsqueeze(0).to(device))

        optimizer.zero_grad()
        train_loss.backward()
        optimizer.step()

    net.eval()
    prediction = torch.max(net(X_test.to(device)), 1)[1]
    accuracy = accuracy_score(y_test,cpu().numpy(), prediction.cpu().numpy())
    print(f&quot;Epoch {epoch + 1} test accuracy is {accuracy}.&quot;)

then you'll notice that it only takes two epochs for the model to reach 96% accuracy.
