Performance metrics changing significantly based on batch size

Question

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

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

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.

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:

• Tune (reduce) the learning rate. You can check the answers to this SO question for some heuristics on how to choose the value. Apart from that, I recommend the article Don't Decay the Learning Rate, Increase the Batch Size.

• Use an adaptive optimizer, like Adam.