Hi,
I am trying to retrain a 3D CNN model from a research article and I run into overfitting issues even upon implementing data augmentation on the fly to avoid overfitting.
I can see that my model learns and then starts to oscillate along the same loss numbers.
Any suggestions on how to improve or how I should proceed in preventing the model from overfitting will be of great help. For further reference, I am posting my code below and the loss. Thanks
What I tried to avoid overfitting:

image augmentation,
regularization.
Dropout layers
Increasing the batch size
Shuffling the data-loader

Training dataset size: approx 358 images+augmentation(every epoch)
class T_LEAP(torch.nn.Module):
    """T_LEAP ARCHITECTURE"""

    def __init__(self):
        super(T_LEAP, self).__init__()
        self.encoder = torch.nn.Sequential(
                                          torch.nn.Conv3d(3, 64, kernel_size=(3,3,3), stride=(1,1,1), padding=(1,1,1)),
                                          torch.nn.BatchNorm3d(num_features=64),
                                          torch.nn.ReLU(inplace=True),
                                          torch.nn.MaxPool3d(kernel_size=(1,2,2),stride=(1,2,2)),
                                          
                                          
                                           
                                          torch.nn.Conv3d(64, 128, kernel_size=(3,3,3), stride=(1,1,1), padding=(1,1,1)),                                         
                                          torch.nn.BatchNorm3d(num_features=128),
                                          torch.nn.ReLU(inplace=True),
                                          torch.nn.MaxPool3d(kernel_size=(2,2,2),stride=(2,2,2)),
                                          
                                           
                                          torch.nn.Conv3d(128, 256, kernel_size=(3,3,3), stride=(1,1,1), padding=(1,1,1)),
                                          torch.nn.BatchNorm3d(num_features=256),
                                          torch.nn.ReLU(inplace=True),
                                          torch.nn.MaxPool3d(kernel_size=(1,2,2),stride=(1,2,2)),
                                         
                                           
                                          torch.nn.Conv3d(256, 512, kernel_size=(1,3,3), stride=(1,1,1), padding=(1,1,1)),
                                          torch.nn.BatchNorm3d(num_features=512),
                                          torch.nn.ReLU(inplace=True),
                                          
                                          
                                          )
        
        self.decoder = torch.nn.Sequential(
                                          torch.nn.ConvTranspose3d(512, 256, kernel_size=(1,3,3), stride = (1,2,2), padding=(1,1,1), output_padding=(0,1,1)),
                                          torch.nn.BatchNorm3d(num_features=256), 
                                          torch.nn.ReLU(inplace=True),
                                     
                                          torch.nn.Conv3d(256, 256, kernel_size=(1,3,3), stride=(1,1,1), padding=(1,1,1)),
                                          torch.nn.BatchNorm3d(num_features=256),
                                          torch.nn.ReLU(inplace=True),
        
                                          torch.nn.ConvTranspose3d(256,128, kernel_size=(1,3,3), stride = (1,2,2), padding=(1,1,1),output_padding=(0,1,1)), 
                                          torch.nn.BatchNorm3d(num_features=128),
                                          torch.nn.ReLU(inplace=True),
                                          
                                          torch.nn.Conv3d(128, 128, kernel_size=(1,3,3), stride=(1,1,1), padding=(1,1,1)),
                                          torch.nn.BatchNorm3d(num_features=128),
                                          torch.nn.ReLU(inplace=True),
                                          
                                          torch.nn.ConvTranspose3d(128,17, kernel_size=(1,3,3), stride = (1,2,2), padding=(1,1,1),output_padding=(0,1,1)),
                                          torch.nn.BatchNorm3d(num_features=17),
                                          torch.nn.ReLU(inplace=True),
                                          )           
               
    def forward(self, image):
        """PUTTING THE MODEL TOGETHER"""
        encoder = self.encoder(image)
        decoder = self.decoder(encoder)
        out     = torch.nn.Softmax(dim=1)(decoder)
        return out


model = T_LEAP()
if torch.cuda.is_available():
  input = torch.rand(1,3, 2, 200, 200).cuda()
  model = model.cuda()
  print(model(input).shape)
else:
  summary(model,
          input_size=(3, 2, 200, 200),
          batch_size=1
          )

Training loss:

other stuff i checked:

image434×657 11.1 KB

Training loop:
DEVICE     = "cuda" if torch.cuda.is_available() else "cpu"
EPOCHS     = 35
lr         = 1e-3
model      = T_LEAP()
optimizer  = torch.optim.Adam(model.parameters(), lr=lr)
model = model.to(DEVICE)
criterion  = torch.nn.MSELoss()

train_loss = []
for epoch in range(EPOCHS): 
    running_loss = 0.0
    for i, data in enumerate(train_loader, 0):
        # Get[inputs, labels]
        inputs, labels = data["images"], data["heatmaps"]

        inputs = inputs.permute(0,2,1,3,4).to(DEVICE,dtype=torch.float)
        labels = labels.unsqueeze(0).permute(1,4,0,2,3).to(DEVICE,dtype=torch.float)
        
        # zero the parameter gradients
        optimizer.zero_grad()

        # forward + backward + optimize
        outputs = model(inputs)
        loss = criterion(outputs, labels)
        loss.backward()
        optimizer.step()

        running_loss += loss.item()

    print("loss:",running_loss/len(train_loader))
    train_loss.append(running_loss)
        
print('Finished Training')