Is my code correct for calculating score during validation?

username: jawi289p
cooked: <p>Can anyone please check my code for calculating IoU (0 &amp; 1 classes) and mIoU scores during the validation? Is this correct?</p>
<pre><code class="lang-auto">train_losses = []
val_losses = []

# Function to compute IoU
def compute_iou_binary_torch(y_true, y_pred):
    """
    Computes IoU for binary classification directly from PyTorch tensors.
    
    Args:
        y_true (tensor): Ground truth binary tensor (0s and 1s).
        y_pred (tensor): Predicted binary tensor (0s and 1s).
    
    Returns:
        IoU (list): Intersection over Union for each class.
        mIoU (float): Mean IoU across classes.
    """
    # Convert tensors to NumPy arrays
    y_true = y_true.cpu().numpy().flatten()
    y_pred = y_pred.cpu().numpy().flatten()
    
    # Compute confusion matrix
    current = confusion_matrix(y_true, y_pred, labels=[1, 0])
    intersection = np.diag(current)
    ground_truth_set = current.sum(axis=1)
    predicted_set = current.sum(axis=0)
    union = (ground_truth_set + predicted_set - intersection).astype(np.float32)
    
    with np.errstate(divide='ignore', invalid='ignore'):
        IoU = np.divide(intersection, union)
        IoU[np.isnan(IoU)] = 0  # Replace NaN values with 0
    
    mIoU = np.nanmean(IoU)
    return IoU, mIoU

# Training and validation loop
for epoch in range(n_eps):
    model.train()
    train_loss = 0.0

    # Training loop
    for images, masks in tqdm(train_loader):
        images, masks = images.to(device), masks.to(device)

        optimizer.zero_grad()
        outputs = model(images)

        loss = criterion(outputs, masks)
        loss.backward()
        optimizer.step()

        train_loss += loss.item()

    avg_train_loss = train_loss / len(train_loader)
    train_losses.append(avg_train_loss)
    print(f"Epoch [{epoch+1}/{n_eps}], Train Loss: {avg_train_loss:.4f}")

    model.eval()
    val_loss = 0.0
    iou_scores = []  # Reset IoU scores at the start of the validation loop
    class_iou_scores = []  # Reset class-wise IoU scores at the start of the validation loop

    # Validation loop
    with torch.no_grad():
        for images, masks in val_loader:
            images, masks = images.to(device), masks.to(device)
            outputs = model(images)

            # Calculate loss on raw outputs
            val_loss += criterion(outputs, masks).item()

            # Apply sigmoid and thresholding for binary prediction
            predictions = (outputs.sigmoid() &gt; 0.5).float()

            # Compute IoU for the current batch
            batch_iou, batch_miou = compute_iou_binary_torch(masks, predictions)
            class_iou_scores.append(batch_iou)
            iou_scores.append(batch_miou)

    avg_val_loss = val_loss / len(val_loader)
    avg_iou = np.mean(iou_scores)  # Compute average IoU across all batches
    avg_class_iou = np.mean(class_iou_scores, axis=0)  # Compute average IoU for each class
    val_losses.append(avg_val_loss)

    print(f"Epoch [{epoch+1}/{n_eps}], Val Loss: {avg_val_loss:.4f}")
    print(f"Epoch [{epoch+1}/{n_eps}], Class-wise IoU: {avg_class_iou}")
    print(f"Epoch [{epoch+1}/{n_eps}], Mean IoU: {avg_iou:.4f}")
</code></pre>

username: jawi289p
cooked: <p><a class="mention" href="/u/ptrblck">@ptrblck</a> can you please check?</p>