| import torch |
| from torchmetrics.classification import MulticlassAccuracy, MulticlassAveragePrecision |
|
|
|
|
| class MetricTracker: |
| def __init__(self, num_classes, device): |
| self.num_classes = num_classes |
| self.device = device |
|
|
| self.map_metric = MulticlassAveragePrecision(num_classes=num_classes).to(device) |
| self.acc_metric = MulticlassAccuracy(num_classes=num_classes).to(device) |
|
|
| self.reset() |
|
|
| def reset(self): |
| self.map_metric.reset() |
| self.acc_metric.reset() |
| self.loss_sum = 0 |
| self.count = 0 |
|
|
| def update(self, preds, targets, loss=None, skip_metrics=False): |
| """ |
| preds: logits [B, C] |
| targets: [B] or soft labels [B, C] |
| skip_metrics: If True, only loss is tracked. Use for MixUp/CutMix batches. |
| """ |
| if targets.ndim > 1: |
| hard_targets = targets.argmax(dim=1) |
| else: |
| hard_targets = targets |
| if not skip_metrics: |
| self.map_metric.update(preds, hard_targets) |
| self.acc_metric.update(preds, hard_targets) |
|
|
| if loss is not None: |
| self.loss_sum += loss * preds.size(0) |
| self.count += preds.size(0) |
|
|
| def compute(self): |
| mAP = self.map_metric.compute().item() |
| acc = self.acc_metric.compute().item() |
| avg_loss = self.loss_sum / max(self.count, 1) |
| return {"mAP": mAP, "accuracy": acc, "loss": avg_loss} |
|
|
