CNN model for RGB images giving 0% accuracy
I am trying to train a CNN model on CelebA (RGB images) dataset. But, when I train the model and check its accuracy it is 0% or close to 0%. I think the issue is in the ConNeuralNet function or the hyperparameters but due to my limited knowledge I'm not sure what I'm missing here. Can someone please help. Thanks
# Creating a simple network
class ConvNeuralNet(torch.nn.Module):          
 
    def __init__(self, num_classes=10178):     
        super(ConvNeuralNet, self).__init__()
        self.conv_layer1 = nn.Conv2d(in_channels=3, out_channels=32, kernel_size=3)
        self.conv_layer2 = nn.Conv2d(in_channels=32, out_channels=32, kernel_size=3)
        self.max_pool1 = nn.MaxPool2d(kernel_size = 2, stride = 2)
    
        self.conv_layer3 = nn.Conv2d(in_channels=32, out_channels=64, kernel_size=3)
        self.conv_layer4 = nn.Conv2d(in_channels=64, out_channels=64, kernel_size=3)
        self.max_pool2 = nn.MaxPool2d(kernel_size = 2, stride = 2)
    
        self.fc1 = nn.Linear(13312, 128)
        self.relu1 = nn.ReLU()
        self.fc2 = nn.Linear(128, num_classes)
    
    def forward(self, x):
        out = self.conv_layer1(x)
        out = self.conv_layer2(out)
        out = self.max_pool1(out)
    
        out = self.conv_layer3(out)
        out = self.conv_layer4(out)
        out = self.max_pool2(out)
            
        out = out.reshape(out.size(0), -1)
    
        out = self.fc1(out)
        out = self.relu1(out)
        out = self.fc2(out)
        return F.log_softmax(out,dim=-1)

def trainTorch(torch_model, train_loader, test_loader,
    nb_epochs=NB_EPOCHS, batch_size=BATCH_SIZE, train_end=-1, test_end=-1, learning_rate=LEARNING_RATE, optimizer=None):

    train_loss = []
    total = 0
    correct = 0
    step = 0
    for _epoch in range(nb_epochs):
      for xs, ys in train_loader:
        xs, ys = Variable(xs), Variable(ys)
        if torch.cuda.is_available():
          xs, ys = xs.cuda(), ys.cuda()
      
      
        optimizer.zero_grad()
        preds = torch_model(xs)
        preds = F.log_softmax(preds, dim=1)
        loss = F.cross_entropy(preds, ys)
        loss.backward()
        train_loss.append(loss.data.item())
        optimizer.step()  # update gradients

        preds_np = preds.cpu().detach().numpy()
        correct += (np.argmax(preds_np, axis=1) == ys.cpu().detach().numpy()).sum()
        total += train_loader.batch_size
        step += 1
        if total % 1000 == 0:
          acc = float(correct) / total
          print('[%s] Training accuracy: %.2f%%' % (step, acc * 100))
          total = 0
          correct = 0


nb_epochs = 8
image_size = 64
batch_size = 64
num_classes = 10178
learning_rate = 0.001
num_epochs = 8
# Device will determine whether to run the training on GPU or CPU.
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')


trans = transforms.Compose([
    transforms.Resize(image_size),
    transforms.ToTensor(),
    transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))
])


train_loader = torch.utils.data.DataLoader(
    datasets.CelebA('data', split='train', target_type='identity', transform=trans,     download=&quot;True&quot;),
    batch_size=batch_size, shuffle=True)

test_loader = torch.utils.data.DataLoader(
    datasets.CelebA('data', split='test', target_type='identity', transform=trans),
    batch_size=batch_size)


#Training the model
print(&quot;Training Model&quot;)
# Set optimizer with optimizer
optimizer = torch.optim.SGD(model1.parameters(), lr=learning_rate, weight_decay =     0.005, momentum = 0.9)
total_step = len(train_loader)
trainTorch(model1, train_loader, test_loader, nb_epochs, batch_size, train_end, test_end, learning_rate, optimizer = optimizer)


**Update I ran the code for a bit to see if it would start converging. One thing is that there are over 10,000 classes. With a batch size of 64 this means that it will take more than 150 mini-batches before your model has seen every class in your dataset. You certanly shouldn't expect the model to start achieving accurate predictions within a few hundred steps.
When I printed the loss value I noticed it was decreasing very slowly. I changed to learning rate to 0.01 and it started decreasing faster.
Also, your model is very shallow for a face recognition model. You're better off using something like a resnet variant (e.g. resnet-50 or resnet-101 from torchvision), rather than custom rolling your own model.
Primary changes include

Learning rate increased
Fix the loss function
Remove log_softmax from output of model
Add activation to the conv layers

IMO the comments about softmax are a bit misleading since you don't need to softmax the output of your model if you are using cross_entropy. You also don't need softmax to get the argmax of the prediction since both softmax and log_softmax don't change the relative ordering of the predictions (i.e. both softmax and log are strictly increasing functions).
IMO the comment about using average pooling to reduce the input size of the first fc layer is a good one and may improve performance, but you'll need to experiment with that one to find good parameters for it so I left it out of this answer.
import torch
import torch.nn as nn
import torch.nn.functional as F
import numpy as np
from torchvision import datasets, transforms

# Creating a simple network
class ConvNeuralNet(torch.nn.Module):          
 
    def __init__(self, num_classes=10178):     
        super(ConvNeuralNet, self).__init__()
        self.conv_layer1 = nn.Conv2d(in_channels=3, out_channels=32, kernel_size=3)
        self.conv_layer2 = nn.Conv2d(in_channels=32, out_channels=32, kernel_size=3)
        self.max_pool1 = nn.MaxPool2d(kernel_size = 2, stride = 2)
    
        self.conv_layer3 = nn.Conv2d(in_channels=32, out_channels=64, kernel_size=3)
        self.conv_layer4 = nn.Conv2d(in_channels=64, out_channels=64, kernel_size=3)
        self.max_pool2 = nn.MaxPool2d(kernel_size = 2, stride = 2)
    
        self.fc1 = nn.Linear(13312, 128)
        self.relu1 = nn.ReLU()
        self.fc2 = nn.Linear(128, num_classes)
    
    def forward(self, x):
        # note the relu activations on the conv layers
        out = F.relu(self.conv_layer1(x))
        out = F.relu(self.conv_layer2(out))
        out = self.max_pool1(out)
    
        out = F.relu(self.conv_layer3(out))
        out = F.relu(self.conv_layer4(out))
        out = self.max_pool2(out)
        
        # you may want an adaptive average pool 2d here to reduce size of feature map further

        out = out.reshape(out.size(0), -1)
    
        out = self.fc1(out)
        out = self.relu1(out)
        out = self.fc2(out)
        # return raw logits, not log-softmax output
        return out

def trainTorch(torch_model, train_loader, test_loader, nb_epochs, batch_size, learning_rate, optimizer):

    train_loss = []
    total = 0
    correct = 0
    step = 0
    for _epoch in range(nb_epochs):
      for xs, ys in train_loader:
        # the Variable interface has been deprecated for years, it is effectively a no-op in modern pytorch
        # see: https://pytorch.org/docs/stable/autograd.html#variable-deprecated
        if torch.cuda.is_available():
          xs, ys = xs.cuda(), ys.cuda()

        optimizer.zero_grad()
        logits = torch_model(xs)
        # don't softmax or log-softmax the inputs to cross_entropy
        loss = F.cross_entropy(logits, ys)
        # The following is equivalent but less numerically stable
        # loss = F.nll_loss(F.log_softmax(logits), ys)
        loss.backward()
        train_loss.append(loss.item())
        optimizer.step()  # update gradients

        logits_np = logits.cpu().detach().numpy()
        correct += (np.argmax(logits_np, axis=1) == ys.cpu().detach().numpy()).sum()
        total += train_loader.batch_size
        step += 1
        if step % 200 == 0:
          acc = float(correct) / total
          avg_loss = sum(train_loss) / len(train_loss)
          print(f'[{step}] Training accuracy: {acc*100:.2f}%     Training loss: {avg_loss:.4f}')
          total = 0
          correct = 0
          train_loss = []


nb_epochs = 8
image_size = 64
batch_size = 64
num_classes = 10178
# increased learning rate to 0.01
learning_rate = 0.01
num_epochs = 8
# Device will determine whether to run the training on GPU or CPU.
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')

trans = transforms.Compose([
    transforms.Resize(image_size),
    transforms.ToTensor(),
    transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))
])

train_loader = torch.utils.data.DataLoader(
    datasets.CelebA('data', split='train', target_type='identity', transform=trans, download=True),
    batch_size=batch_size, shuffle=True)

test_loader = torch.utils.data.DataLoader(
    datasets.CelebA('data', split='test', target_type='identity', transform=trans),
    batch_size=batch_size)

model = ConvNeuralNet(num_classes)
if torch.cuda.is_available():
    model.cuda()

#Training the model
print(&quot;Training Model&quot;)
# Set optimizer with optimizer
optimizer = torch.optim.SGD(model.parameters(), lr=learning_rate, weight_decay=0.005, momentum=0.9)
total_step = len(train_loader)
trainTorch(model, train_loader, test_loader, nb_epochs, batch_size, learning_rate, optimizer=optimizer)

Output
Training Model
[200] Training accuracy: 0.00%     Training loss: 9.2286
[400] Training accuracy: 0.02%     Training loss: 9.2286
[600] Training accuracy: 0.04%     Training loss: 9.2265
[800] Training accuracy: 0.00%     Training loss: 9.2253
[1000] Training accuracy: 0.00%     Training loss: 9.2222
[1200] Training accuracy: 0.00%     Training loss: 9.2105
[1400] Training accuracy: 0.02%     Training loss: 9.1776
[1600] Training accuracy: 0.03%     Training loss: 9.1329
[1800] Training accuracy: 0.02%     Training loss: 9.1013
[2000] Training accuracy: 0.02%     Training loss: 9.0830
[2200] Training accuracy: 0.02%     Training loss: 9.0715
[2400] Training accuracy: 0.01%     Training loss: 9.0622
[2600] Training accuracy: 0.02%     Training loss: 9.0456
[2800] Training accuracy: 0.00%     Training loss: 9.0301
[3000] Training accuracy: 0.00%     Training loss: 9.0357
[3200] Training accuracy: 0.02%     Training loss: 9.0402
[3400] Training accuracy: 0.02%     Training loss: 9.0321
[3600] Training accuracy: 0.02%     Training loss: 9.0217
[3800] Training accuracy: 0.02%     Training loss: 8.9757
[4000] Training accuracy: 0.09%     Training loss: 8.9059
[4200] Training accuracy: 0.09%     Training loss: 8.8331
[4400] Training accuracy: 0.09%     Training loss: 8.7601
[4600] Training accuracy: 0.09%     Training loss: 8.7356
[4800] Training accuracy: 0.10%     Training loss: 8.6717
[5000] Training accuracy: 0.12%     Training loss: 8.6311
[5200] Training accuracy: 0.16%     Training loss: 8.5515
[5400] Training accuracy: 0.16%     Training loss: 8.4943
[5600] Training accuracy: 0.14%     Training loss: 8.4345
[5800] Training accuracy: 0.14%     Training loss: 8.4107
[6000] Training accuracy: 0.18%     Training loss: 8.3317
[6200] Training accuracy: 0.22%     Training loss: 8.2716
[6400] Training accuracy: 0.31%     Training loss: 8.1934
[6600] Training accuracy: 0.30%     Training loss: 8.1500
[6800] Training accuracy: 0.35%     Training loss: 8.0979
[7000] Training accuracy: 0.21%     Training loss: 8.0739
[7200] Training accuracy: 0.44%     Training loss: 8.0220
[7400] Training accuracy: 0.29%     Training loss: 7.9819

From the output we see the loss is decreasing and the accuracy is starting to increase. Its hard to predict how well this will work and when it will converge but this is a good start. You'll probably need to use a better model and a learning rate scheduler to get better performance.
For example, just switching for a resnet-50
model = torchvision.models.resnet50(pretrained=True)
model.fc = nn.Linear(model.fc.in_features, num_classes)

The model starts converging much faster
Training Model
[200] Training accuracy: 0.05%     Training loss: 9.1942
[400] Training accuracy: 0.05%     Training loss: 8.9244
[600] Training accuracy: 0.15%     Training loss: 8.5936
[800] Training accuracy: 0.30%     Training loss: 8.3147
[1000] Training accuracy: 0.39%     Training loss: 8.0745
[1200] Training accuracy: 0.43%     Training loss: 7.9146
[1400] Training accuracy: 0.45%     Training loss: 7.7706
[1600] Training accuracy: 0.64%     Training loss: 7.6551
[1800] Training accuracy: 0.68%     Training loss: 7.5784
[2000] Training accuracy: 0.74%     Training loss: 7.5327
[2200] Training accuracy: 0.72%     Training loss: 7.4689
[2400] Training accuracy: 0.63%     Training loss: 7.4378
[2600] Training accuracy: 0.83%     Training loss: 7.3789
[2800] Training accuracy: 0.90%     Training loss: 7.2812
[3000] Training accuracy: 0.84%     Training loss: 7.2771
[3200] Training accuracy: 0.96%     Training loss: 7.2536
[3400] Training accuracy: 1.00%     Training loss: 7.2538

