I am working on image classification problem with 3 classes. There are some issues with my dataset viz, the first class is quite unbalanced ,i.,e the number of images in my whole corpus is in this proportion - 201:1000:1000 . Hence, in the training dataset which is 80% (1760 images), I used weighted random sampling to oversample the 1st class and uses some image augmentation techniques(rotation, horizontal flip, color jitter). Then I fit the Resnet-50 model and Resnet -18 model and only train the last layer. After 10 epochs, (batch size -16 and learning rate-0.0001 ), the loss and accuracy plot on Resnet-50 model are shown below.
My question is should I train for more epoch, there is quite overfitting here. How to solve this problem?