This is the training code I’m using. I had to create my own Resize function since I’m dealign with TIFF images. I use F.interpolate for that.
I’ve been testing the network with only 1 image and I can’t overfit. I also implemented different models from the web and all of them behave in the same way, that’s why I think the problem is on my training, data loader or transformation code.
What do you think could be the problem?
import torch
import unet3D
import numpy as np
import transformations_3_D
import torchvision.transforms as transforms
import torch.nn as nn
import torch.optim as optim
import visdom
import time
from sklearn.metrics import jaccard_score
from torch.functional import F
from dataset_3D import MacularHole
from torch.utils import data
# from images3D_visualizer import ImageVisualizer
from matplotlib import pyplot as plt
import unet3d_github

# Resize function

# Resize function
def labels_resize(tensor, size, normalize=False):
    tensor = F.interpolate(tensor, size=size, mode="nearest")
    if normalize:
        if torch.min(tensor) == 0 and torch.max(tensor) == 0:
            tensor = tensor
        else:
            tensor = (tensor - torch.min(tensor)) / (torch.max(tensor) - torch.min(tensor))

    return tensor


def prediction_resize(tensor, size):
    tensor = F.interpolate(tensor, size=size, mode="trilinear", align_corners=True)
    return tensor


def main():
    device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")

    transform = transforms.Compose([transformations_3_D.Resize((116, 132, 132))])

    seed = np.random.randint(0, 1000)

    train_set = MacularHole("dataset_first_tests_45_imgs/", mode='train', seed=seed, transform=transform,
                            training_size=1, valid_size=9, testing_size=9)

    train_loader = torch.utils.data.DataLoader(train_set, batch_size=1, shuffle=True, num_workers=16)
  
    # model = unet3D.UNet()

    model = model.to(device)

    criterion = nn.BCELoss()
    optimizer = optim.Adam(model.parameters(), lr=0.0001)

    # Lists of data for visdom
    max_train_predictions = []
    train_performance_values = []
    max_valid_predictions = []
    valid_performance_values = []
    train_loss_values = []
    valid_loss_values = []
    test_performance_values = []

    epochs = 5000
    valid_performance_min = 0.0
    train_max_value = 0.0
    valid_max_value = 0.0
    t0 = time.time()
    for i in range(epochs):
        train_loss = 0.0
        valid_loss = 0.0
        jaccard_scores_train = []

        model.train()
        for batch, (images, labels) in enumerate(train_loader):
            images, labels = images.float(), labels.float()
            images, labels = images.to(device), labels.to(device)
            labels = labels_resize(labels, (49, 200, 200))

            optimizer.zero_grad()
            train_predictions = model(images)
            resize_labels = labels_resize(labels, (28, 44, 44), normalize=True)
            train_max_value = torch.max(train_predictions)
            loss = criterion(train_predictions, resize_labels)
            loss.backward()
            optimizer.step()
            train_loss += loss.item()

            # jaccard scores
            with torch.no_grad():

                train_predictions = prediction_resize(train_predictions, (49, 200, 200))
                # Binarizing labels
                labels = labels != 0
                # Thresholding predictions
                binary_train_predictions = train_predictions > 0.5
                labels = labels.detach().cpu().numpy().astype(int)
                binary_train_predictions = binary_train_predictions.detach().cpu().numpy().astype(int)
                # Flattening numoy arrays for jaccard_score and adding each value to a list
                jaccard_scores_train += [jaccard_score(labels.flatten(), binary_train_predictions.flatten())]

        # calculating loss and performance in each epoch
        train_loss = train_loss / len(train_set)
        train_loss_values.append(train_loss)
        train_performance = np.mean(jaccard_scores_train)
        train_performance_values.append(train_performance)

        print(f'Epoch: {i + 1} |' f'\tTrain loss: {train_loss} |' f'\t| Train performance: {train_performance} |' 
              f'\t| Max training: {train_max_value} |')



if __name__ == '__main__':
    # vis = visdom.Visdom()
    main()