I am working on a binary classification problem where there is significant class imbalance (minority class makes up nearly 10%). The dataset has ~15,000 observations and I have split this in to a training, validation and test set (that are stratified).

Using PyTorch I build a neural network with 5 fully connected layers (using ReLU activation), CrossEntropyLoss and SGD optimiser. Below is parts of my code

The problem is that my training vs validation loss changes a lot based on batch size (passed in the DataLoader). If I use a batch size of 64, the loss functions look like

enter image description here

which is quite odd. But if I use an unconventionally large batch size of say 1000, it looks like: enter image description here

This looks more familiar but I can't make sense of what is going wrong here. I am also seeing that the training set reaches a high recall fairly quickly (after ~4 epochs) while the validation set improves slowly. So there seems to be an issue of overfitting as well.

I don't really know where I am going wrong: my neural network architecture consists of 5 fully connected layers with appropriate input and output dimensions. I initialise the weights. The forward function applies ReLU to the inputs (I don't use Softmax because I only need to classify 0 or 1 so I thought I can simply use argmax, see the 'c' variable in the code above).

I have tried setting Shuffle to true in the training_loader but this produces highly fluctuating training loss values.


1 Answer 1


Batch size is very related to the learning rate, especially in non-adaptive optimizers like the vanilla SGD that you are using.

I would suggest two alternatives:

  • $\begingroup$ That did it! I changed the learning rate from 0.01 to 0.001 and even 0.0001 and the performance metrics look very reasonable. The same thing when changing to Adam. While I understand the logic behind SGD, Adam is not as clear to me but I will play around with the parameters and try to fine-tune them. Many thanks for the help :) $\endgroup$
    – BenBernke
    Commented Apr 16, 2021 at 17:29

Your Answer

By clicking “Post Your Answer”, you agree to our terms of service and acknowledge you have read our privacy policy.

Not the answer you're looking for? Browse other questions tagged or ask your own question.