Correct way to calculate IoU for each class

jawi289p: <p>I am doing multi class segmentation and I want to know what is the correct way for calculating and displaying iou for each class during the validation of the data. Any help will be appreciated. Thank you</p>
<p>Below is my code:</p>
<p>Validation</p>
<pre><code class="lang-auto">    model.eval()
    running_val_loss = 0.0
    with torch.no_grad():
        for images, masks in val_loader:  # Replace val_loader with your DataLoader
            images, masks = images.to(device), masks.to(device)

            outputs = model(images)

            # Calculate validation loss
            val_loss = criterion(outputs, masks)
            running_val_loss += val_loss.item()

            # Call the function to calculate mIoU and IoU metrics
            miou = calculate_iou_scores(outputs, masks, num_classes=5)


    avg_val_loss = running_val_loss / len(val_loader)
    val_losses.append(avg_val_loss)

    # Print validation results
    print(f'Epoch [{epoch+1}/{n_eps}], Val Loss: {avg_val_loss:.4f}, mIoU: {miou:.4f}')
</code></pre>
<p>calculate_iou_score function</p>
<pre><code class="lang-auto">def calculate_iou_scores(outputs, targets, num_classes):
    # Convert logits to class predictions
    outputs = torch.argmax(outputs, dim=1)

    # Initialize IoU metric
    iou = torchmetrics.JaccardIndex(num_classes=num_classes, average='macro', task='multiclass').to(device)


    # Calculate IoU and mIoU
    iou.update(outputs, targets)
    iou_score = iou.compute()

    return iou_score.item()
</code></pre>

Soumya_Kundu: <p>Something like this may work:</p>
<pre><code class="lang-auto"># Val
model.eval()
running_val_loss = 0.0
class_iou_sums = [0] * num_classes
with torch.no_grad():
    for images, masks in val_loader:
        outputs = model(images)
        val_loss = criterion(outputs, masks)
        running_val_loss += val_loss.item()
        iou_scores, batch_miou = calculate_iou_scores(outputs, masks, num_classes=num_classes)
        for i, iou in enumerate(iou_scores):
            class_iou_sums[i] += iou

avg_val_loss = running_val_loss / len(val_loader)
class_ious = [iou_sum / len(val_loader) for iou_sum in class_iou_sums]
miou = sum(class_ious) / num_classes

# Iterative Results
print(f'Epoch [1/1], Val Loss: {avg_val_loss:.4f}, mIoU: {miou:.4f}')
print('IoU for each class:')
for i, iou in enumerate(class_ious):
    print(f'Class {i}: {iou:.4f}')
</code></pre>

jawi289p: <p>this code is not working…!</p>

Soumya_Kundu: <p>It is on my machine (haha)</p>
<p>Also, do you mean not working as in logically or its throwing some error?</p>
<pre><code class="lang-auto"># pip install torchmetrics
import torch
import torchmetrics

# Dummy values
num_classes = 5
num_batches = 100
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

# Dummy model (just for testing)
class DummyModel(torch.nn.Module):
    def forward(self, x):
        return torch.randn(x.size(0), num_classes, 16, 16)

model = DummyModel().to(device)

criterion = torch.nn.CrossEntropyLoss()

# Dummy validation loader
class DummyLoader:
    def __init__(self, num_batches):
        self.num_batches = num_batches

    def __iter__(self):
        for _ in range(self.num_batches):
            images = torch.randn(4, 3, 16, 16).to(device)
            masks = torch.randint(0, num_classes, (4, 16, 16)).to(device)
            yield images, masks

    def __len__(self):
        return self.num_batches

val_loader = DummyLoader(num_batches)

def calculate_iou_scores(outputs, targets, num_classes):
    outputs = torch.argmax(outputs, dim=1)
    iou_metrics = [torchmetrics.JaccardIndex(num_classes=num_classes, average='none', task='multiclass').to(device)
                   for _ in range(num_classes)]
    iou_scores = []
    for class_idx in range(num_classes):
        iou_metrics[class_idx].update(outputs, targets)
        iou_score = iou_metrics[class_idx].compute()
        iou_scores.append(iou_score[class_idx].item())
    miou = sum(iou_scores) / num_classes
    return iou_scores, miou

# Val loop
model.eval()
running_val_loss = 0.0
class_iou_sums = [0] * num_classes
with torch.no_grad():
    for images, masks in val_loader:
        outputs = model(images)
        val_loss = criterion(outputs, masks)
        running_val_loss += val_loss.item()
        iou_scores, batch_miou = calculate_iou_scores(outputs, masks, num_classes=num_classes)
        for i, iou in enumerate(iou_scores):
            class_iou_sums[i] += iou

avg_val_loss = running_val_loss / len(val_loader)
class_ious = [iou_sum / len(val_loader) for iou_sum in class_iou_sums]
miou = sum(class_ious) / num_classes

# Print validation results
print(f'Epoch [1/1], Val Loss: {avg_val_loss:.4f}, mIoU: {miou:.4f}')
print('IoU for each class:')
for i, iou in enumerate(class_ious):
    print(f'Class {i}: {iou:.4f}')
</code></pre>
<p>Maybe try running this.</p>

jawi289p: <p><a class="mention" href="/u/soumya_kundu">@Soumya_Kundu</a> thanks man it is working<br>
previously there was issue in the calculate_iou_scores function</p>