Train loss decreases and validation loss increases (Overfitting), What Can I do?

Question

I'm training my network (UNET) on Oxford Pets Dataset. Looking at the plots of train loss and validation loss, I'm pretty sure that the network overfits.

I trained the network for 100 epochs, with a learning rate of 0,0001 and a batch size=1.

My question is: Could it be since I have used a batch size=1? If I use a batch size higher, for example, a batch size = 8, then the network at each epoch should move the weights based on 8 images, is it right? Could it happen because I have trained too much my network (many steps)?

Additionally, a strange behavior that I have seen is that at the start the network seems to learn something about some images, but then It gets lost.

Could anyone give me some suggestions on how I could avoid this problem?

Answer 1

A batch size of 1 is extreme. Something like 32 or so is much more typical. Unless you run into memory issues, I typically set 32 as my initial batch size (when you try to go larger problems also arise, there's a whole literature on training with large batch sizes to avoid those). Much smaller (certainly once we get into single digits, but sometimes even before that) tends to cause problems. Some, you can avoid by gradient accumulation, but some not. One of the best known issues with small batch sizes arises when you have BatchNorm layers in your neural network (I'm guessing you have them). The batch size is the first thing, I'd investigate. Also, how did you pick your learning rate (the popular learning rate finder approach?)? Making sure it is set sensibly is important.

The other typical idea to avoid overfitting is data augmentation: a sensible extent of data augmentation (and making sure the augmentation deals properly with the masks - as e.g. the albumentations python library and many other libraries do - by always transforming the mask, too, so it remains in the right location on the image no matter whether you rotate, zoom in or whatever).

Did you try to follow the approach of a good standard example such as the one given in the first chapter of the "Deep Learning for Coders with fastai & PyTorch" book (the example there using a convenient library that does a lot of sensible default things, you can then look into their code for more details; see also the matching course video)?