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
Below is my code:
Validation
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}')
calculate_iou_score function
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()
Soumya_Kundu
August 22, 2024, 11:32am
2
Something like this may work:
# 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}')
this code is not working…!
Soumya_Kundu
August 22, 2024, 1:08pm
4
It is on my machine (haha)
Also, do you mean not working as in logically or its throwing some error?
# 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}')
Maybe try running this.
1 Like
@Soumya_Kundu thanks man it is working