I’m training a model using pytorch on 3d MR brain images.
I have defined a few functions to handle training and evaluation but it seems like the model is not training at all.
The loss changes but the accuracy remains the same. precision and recall metrics are always zero
I can’t figure out what I’m doing wrong here.
My model is based on VGG16 but customized for 3D inputs.
The task is binary classification and there is a sigmoid layer at the end of model pipeline.

training function:

def train_step(model, loss_fn, optimizer, dataloader):
    model.train()

    total_loss = 0
    total_correct = 0
    for (X, Y) in tqdm(dataloader):
        X, Y = X.to(device), Y.to(device)

        logit_output = model(X)
        
        loss = loss_fn(logit_output, Y.view(-1, 1).float())

        optimizer.zero_grad()
        loss.backward()
        optimizer.step()

        total_loss += loss.item()

        predictions = torch.argmax(logit_output, dim=1)
        total_correct += (predictions == Y).sum().item()

    return total_loss, total_correct

evaluation function:

def eval_step(model, loss_fn, dataloader):
    model.eval()

    total_loss = 0
    total_correct = 0
    y_predictions = []
    y_true = []

    with torch.inference_mode():
        for (X, Y) in tqdm(dataloader):
            X, Y = X.to(device), Y.to(device)

            prob_output = model(X)

            loss = loss_fn(prob_output, Y.view(-1, 1).float())
            total_loss += loss.item()

            # Extend lists with raw outputs and true labels for AUC
            y_predictions.extend(prob_output.squeeze(0 if len(prob_output) == 1 else 1).cpu().numpy())
            y_true.extend(Y.cpu().numpy())
    
    # Convert to numpy arrays for metrics calculation
    y_predictions = np.array(y_predictions)
    y_true = np.array(y_true)

    # Calculate AUC
    auc = roc_auc_score(y_true, y_predictions)

    # Calculate precision and recall
    predicted_labels = (y_predictions > 0.3).astype(int)
    precision = precision_score(y_true, predicted_labels)
    recall = recall_score(y_true, predicted_labels)

    return total_loss, auc, precision, recall

model itself:

kleic_model = VGG3D(input_image_size=(1, 16, 256, 256),
               layers=[1, 2, 4],
               convolution_kernel_size=(3, 3, 3),
               pool_size=(2, 2, 2),
               strides=(2, 2, 2),
               number_of_dense_units=128,
               dropout_rate=0.5,
               number_of_outputs=1,
               style=16,
               lowest_resolution=16)

ce_loss = nn.BCELoss()
adam_optim = torch.optim.Adam(params=kleic_model.parameters(), lr=0.001)