[Python] My testing loss isn't improving while I'm using same train and test data

I'm trying to fine tune a model for a segmentation task. To see if everything is working properly I'm trying to make my model overfit on a single image split in 16 different patches.
So in my code, I tried using the exact same dataLoader for my training and testing.
What I don't understand is that my training loss is improving whereas my testing loss isn't moving a single bit. If I'm training and testing on the same data, the results of training loss and testing loss should approximately be the same, right ?

Here is my code:

The customDataSet :

class CustomImageDataset(Dataset):

def __init__(self, seg_dir, img_dir, loc, contrast, transform=None, target_transform=None):
self.img_labels = seg_dir #Set the directory of the images
self.img_dir = img_dir #Set the directory of the labels
self.loc = loc
self.contrast = contrast
self.transform = transform
self.target_transform = target_transform

def __len__(self):
return len([entry for entry in os.listdir(self.img_labels) if os.path.isfile(os.path.join(self.img_labels, entry))])-1

def __getitem__(self, idx):
img_path = os.path.join(self.img_dir, "patch_"+str(idx)+".nii")
image = Image(img_path)
label_path = os.path.join(self.img_labels, "label_"+str(idx)+".nii")
label = Image(label_path)

image_data = image.data
image_data = np.float32(np.expand_dims(image_data, axis = 0))

label_data = label.data
label_data = np.float32(np.expand_dims(label_data, axis = 0))

if self.transform:
image = self.transform(image)
if self.target_transform:
label = self.target_transform(label)

return image_data, label_data


The training and testing loops:

def train_loop(dataloader, model, loss_fn, optimizer, batch_size):
size = len(dataloader.dataset)
model.train()
for batch, (X, y) in enumerate(dataloader):
pred = model(X)
loss = loss_fn(pred, y)

loss.backward()
optimizer.step()
optimizer.zero_grad()
loss, current = loss.item(), batch * batch_size + len(X)
print(f"loss: {loss:>7f} [{current:>5d}/{size:>5d}]")


def test_loop(dataloader, model, loss_fn, batch_size):
# Set the model to evaluation mode - important for batch normalization and dropout layers
model.eval()
size = len(dataloader.dataset)
#print("Size = ", size)
num_batches = len(dataloader)
#print("num_batches = ", num_batches)
test_loss, correct = 0, 0

# Evaluating the model with torch.no_grad() ensures that no gradients are computed during test mode
# also serves to reduce unnecessary gradient computations and memory usage for tensors with requires_grad=True
with torch.no_grad():
for X, y in dataloader:
pred = model(X)
test_loss += loss_fn(pred, y).item()
#print("Test loss = ", test_loss)
correct += (pred.argmax(1) == y).type(torch.float).sum().item()

test_loss /= num_batches
correct /= size
print(f"Test Error: \n Accuracy: {(100*correct):>0.1f}%, Avg loss: {test_loss:>8f} \n")


And my implementation of 3D-DiceLoss : class DiceLoss(nn.Module): def init(self): super(DiceLoss, self).init()

def forward(self, pred, target):
smooth = 0.1
# Calculate intersection and union
intersection = (pred * target).sum()
union = pred.sum() + target.sum()
dice = (2. * intersection) / (union + 1e-8)

#intersection = (pred * target).sum(dim=(2, 3, 4))
#union = pred.sum(dim=(2, 3, 4)) + target.sum(dim=(2, 3, 4))

# Compute Dice Coefficient
#dice = (2. * intersection + smooth) / (union + smooth)

# Return Dice Loss
return 1 - dice


Is all of this an expected behavior due to model.train vs model.eval, or is it something wrong in my code or my entire network ?

Continue reading...