BinaryClassification CNN - validation loss stuck

Hello,

Can some take a look my code and confirm i am doing it right? I am TensorFlow based so i try to recreated TF network which has both training and validation accuracy ~90%

TRAIN_DATA_PATH = "../input/chest-xray-pneumonia/chest_xray/chest_xray/train/"
TEST_DATA_PATH = "../input/chest-xray-pneumonia/chest_xray/chest_xray/test/"
TRANSFORM_IMG = transforms.Compose([
    transforms.Resize((224, 224)),
    transforms.ToTensor(),
    transforms.Normalize([0.485, 0.456, 0.406],[0.229, 0.224, 0.225])
    ])

train_data = torchvision.datasets.ImageFolder(root=TRAIN_DATA_PATH, transform=TRANSFORM_IMG)
train_loader = data.DataLoader(train_data, batch_size=BATCH_SIZE, shuffle=True,  num_workers=4)
test_data = torchvision.datasets.ImageFolder(root=TEST_DATA_PATH, transform=TRANSFORM_IMG)
test_loader  = data.DataLoader(test_data, batch_size=BATCH_SIZE, shuffle=True, num_workers=4) 

Number of train samples: 5216
Number of test samples: 624
Detected Classes are: {‘NORMAL’: 0, ‘PNEUMONIA’: 1}

class CNN(torch.nn.Module): 
    def __init__(self): 
        super().__init__() 
        self.model = torch.nn.Sequential( 
            #Input = 3 x 32 x 32, Output = 16 x 224 x 224 
            torch.nn.Conv2d(in_channels = 3, out_channels = 32, kernel_size = 3, padding = 1), 
            torch.nn.BatchNorm2d(32),
            torch.nn.ReLU(), 
            #Input = 224 x 224 x 224 , Output = 224 x 112 x 112 
            torch.nn.MaxPool2d(kernel_size=2), 
            torch.nn.Dropout(0.2),
            #Input = 224 x 112 x 112, Output = 112 x 112 x 112 
            torch.nn.Conv2d(in_channels = 32, out_channels = 64, kernel_size = 3, padding = 1), 
            torch.nn.BatchNorm2d(64),
            torch.nn.ReLU(), 
            #Input = 112 x 112 x 112, Output = 112 x 56 x 56 
            torch.nn.MaxPool2d(kernel_size=2), 
            torch.nn.Dropout(0.2),
            #Input = 64 x 56 x 56, Output = 64 x28 x 28
            torch.nn.Conv2d(in_channels = 64, out_channels = 128, kernel_size = 3, padding = 1), 
            torch.nn.BatchNorm2d(128),
            torch.nn.ReLU(), 
            #Input = 64 x 28 x 28, Output = 64 x 14 x 14 
            torch.nn.MaxPool2d(kernel_size=2),
            torch.nn.Dropout(0.2),
            torch.nn.Conv2d(in_channels = 128, out_channels = 256, kernel_size = 3, padding = 1),
            torch.nn.BatchNorm2d(256),
            torch.nn.ReLU(), 
            #Input = 64 x 16 x 16, Output = 64 x 16 x 16 
            torch.nn.MaxPool2d(kernel_size=2),
            
            torch.nn.Flatten(), 
            torch.nn.Linear(256*14*14, 64),
            torch.nn.ReLU(),
            torch.nn.Dropout(0.5),
            torch.nn.Linear(64, 1)) 

    def forward(self, x): 
        return self.model(x) 

And most important part, as you can see below i used BCEWithLogitsLoss which has build in sigmoid. So i pass logits not predctions to loss function. I also applied weights bias, and weight_decay to get L2 like in my TF CNN. I am correct?

device = 'cuda' if torch.cuda.is_available() else 'cpu'
model = CNN().to(device) 

num_epochs = 20
learning_rate = 0.001
criterion = torch.nn.BCEWithLogitsLoss(pos_weight = torch.FloatTensor([1.94 / 0.67]).to(device)) 
optimizer = torch.optim.Adam(model.parameters(), lr=learning_rate, weight_decay=0.0001) 
train_losses = []
valid_losses = []
avg_train_losses = []
avg_valid_losses = [] 
train_loss_list = [] 
val_loss_list = [] 
for epoch in range(num_epochs): 
    print(f'Epoch {epoch+1}/{num_epochs}:', end = ' ') 
    
    train_loss = 0
    val_loss = 0

    model.train() 
    for i, (images, labels) in enumerate(train_loader): 

        images = images.to(device) 
        labels = labels.to(device) 

        optimizer.zero_grad() 
        outputs = model(images) 
        loss = criterion(outputs, labels.unsqueeze(1).float())
        
        loss.backward() 
        optimizer.step() 
        train_loss += loss.item() 

    train_loss_list.append(train_loss/len(train_loader)) 
    print(f"Training loss = {train_loss_list[-1]:.5f}")
    
    model.eval()
    for images, labels in test_loader:
        images = images.to(device) 
        labels = labels.to(device)
        output = model(images)
        loss = criterion(output, labels.unsqueeze(1).float())
        val_loss += loss.item() 

    val_loss_list.append(val_loss/len(test_loader)) 
    print(f"Val loss = {val_loss_list[-1]:.5f}")

and i get this (overfitting?):

Training loss = 1.13801
Val loss = 1.72843
Epoch 2/20: Training loss = 0.23643
Val loss = 2.24632
Epoch 3/20: Training loss = 0.22138
Val loss = 1.87088
Epoch 4/20: Training loss = 0.20878
Val loss = 2.09717
Epoch 5/20: Training loss = 0.19411
Val loss = 1.88984
Epoch 6/20: Training loss = 0.17527
Val loss = 1.68840
Epoch 7/20: Training loss = 0.16535
Val loss = 1.75328
Epoch 8/20: Training loss = 0.17526
Val loss = 1.77150
Epoch 9/20: Training loss = 0.16125
Val loss = 1.26466
Epoch 10/20: Training loss = 0.15092
Val loss = 3.07546
Epoch 11/20: Training loss = 0.14819
Val loss = 3.01933
Epoch 12/20: Training loss = 0.14400
Val loss = 3.54311
Epoch 13/20: Training loss = 0.12519
Val loss = 2.37370
Epoch 14/20: Training loss = 0.13651
Val loss = 2.91762
Epoch 15/20: Training loss = 0.12971
Val loss = 3.33134
Epoch 16/20: Training loss = 0.13445
Val loss = 2.27241
Epoch 17/20: Training loss = 0.12127
Val loss = 3.11538
Epoch 18/20: Training loss = 0.13381
Val loss = 2.74372
Epoch 19/20: Training loss = 0.12130
Val loss = 2.09707
Epoch 20/20: Training loss = 0.10609
Val loss = 3.79183

model.eval()
with torch.no_grad():
    correct = 0
    total = 0
    for images, labels in test_loader:
        images = images.to(device)
        labels = labels.to(device)
        outputs = model(images)
        acc = accuracy(outputs, labels.unsqueeze(1).float())

    print(f"Val acc = {acc}")

70%

looks like i need to change archutecture