CNN model is not learn well
I started to learn CNN implementation in PyTorch, and I tried to build CNNs to process the grayscale images with 4 classes from 0 to 3. I got in the beginning accuracy around 0.55. The maximum accuracy I got is ~ 0.683%.
I tried SGD and Adam optimizer with different values for lr and batch_size, but the accuracy is still low.
I used data Augmentation to create more samples, around 4k.
I cannot improve accuracy further and wondered if I could get some advices about what I need to change in CNN structure to increase accuracy.
Loss starts around: Loss: [1.497] then decreases near: Loss: [0.001] then fluctuated up and down around this value.
I spent time reading about similar problems but without luck.
I am using nn.CrossEntropyLoss() for my loss_fn. I don't use softmax for dense layer.
This is the Summary of the CNN model:
-------------------------------------------------------------
        Layer (type)               Output Shape         Param #
=============================================================
            Conv2d-1         [-1, 32, 128, 128]             320
              ReLU-2         [-1, 32, 128, 128]               0
       BatchNorm2d-3         [-1, 32, 128, 128]              64
         MaxPool2d-4           [-1, 32, 64, 64]               0
            Conv2d-5           [-1, 64, 64, 64]          18,496
              ReLU-6           [-1, 64, 64, 64]               0
       BatchNorm2d-7           [-1, 64, 64, 64]             128
         MaxPool2d-8           [-1, 64, 32, 32]               0
            Conv2d-9          [-1, 128, 32, 32]          73,856
             ReLU-10          [-1, 128, 32, 32]               0
      BatchNorm2d-11          [-1, 128, 32, 32]             256
        MaxPool2d-12          [-1, 128, 16, 16]               0
          Flatten-13                [-1, 32768]               0
           Linear-14                  [-1, 512]      16,777,728
             ReLU-15                  [-1, 512]               0
          Dropout-16                  [-1, 512]               0
           Linear-17                    [-1, 4]           2,052
============================================================

I would appreciate the help.

How many images are in the train set ? the test set ? What are the size of the images ? How would you consider the difficulty of classification of the images ? Do you think it should be simple or difficult ?
According to the numbers you have, you're overfitting as your loss is near 0 (meaning nothing much will retropropagate to the weights, i.e your model won't change anymore) and your 68.3% (it's a typo right ?) is from the test set (I suppose). So you don't have any problem to train the network which is a good point.
Then you can search ways of countering overfitting online and here is some &quot;classical&quot; possibilities :
- you may raise the dropout parameter
- putting some regularizer (L1 or L2) to constraint the learning
- early stopping using a validation set
- using a classical and/or lighter convolutional network (resnet,inception) with/without pretrained weights. This latter also depends on your images type (natural, biomedical ...)
- ... a lot more or less difficult to implement
Also technically you are using a softmax layer as it's included in the crossentropyloss of pytorch.
