move metrics from train.py to separate file

Author: DaddyWesker

vla/benchmarks/class_distribution_using_segmentation/metrics.py

@@ -0,0 +1,72 @@
+import torch
+from tqdm import tqdm
+import torch.nn.functional as F
+
+def infer_model(classifier, device, img, loader, model, ttype):
+    if loader.dataset.feature_store is None:
+        feats = model.encode(img.to(device))
+    else:
+        feats = img.to(device)
+    feats = feats.to(ttype)
+    out = classifier(feats)
+    return out
+
+def score_loss(model, classifier, loader, device, ttype):
+    criterion = torch.nn.KLDivLoss(reduction='batchmean')
+    with torch.no_grad():
+        sum_loss = 0
+        for y, img in tqdm(loader):
+            out = infer_model(classifier, device, img, loader, model, ttype)
+            log_probs = F.log_softmax(out, dim=1)
+            loss = criterion(log_probs, y.to(device))
+            sum_loss += loss.item()
+        sum_loss /= len(loader)
+
+    print("Average loss:", sum_loss)
+
+
+def symmetric_topk_mass_recall(logits, targets, k=10):
+
+    probs = F.softmax(logits, dim=-1)
+
+    topk_pred_idx = probs.topk(k, dim=-1).indices
+    target_mass_in_pred_topk = torch.gather(targets, dim=-1, index=topk_pred_idx)
+    recall_p_to_q = target_mass_in_pred_topk.sum(dim=-1)
+
+    topk_target_idx = targets.topk(k, dim=-1).indices
+    pred_mass_in_target_topk = torch.gather(probs, dim=-1, index=topk_target_idx)
+    recall_q_to_p = pred_mass_in_target_topk.sum(dim=-1)
+
+    symmetric_recall = 0.5 * (recall_p_to_q + recall_q_to_p)
+
+    return symmetric_recall.mean()
+
+def symmetric_topk_recall_score(model, classifier, loader, device, ttype):
+    with torch.no_grad():
+        average_recall = 0
+        for y, img in tqdm(loader):
+            out = infer_model(classifier, device, img, loader, model, ttype).squeeze()
+            y = y.squeeze().to(device)
+            recall = symmetric_topk_mass_recall(out, y, k=17)
+            average_recall += recall.item()
+        average_recall /= len(loader)
+    print("Average symmetric top-k recall:", average_recall)
+
+def distribution_iou(logits, targets):
+    probs = F.softmax(logits, dim=-1)
+
+    intersection = torch.minimum(probs, targets).sum(dim=-1)
+    union = torch.maximum(probs, targets).sum(dim=-1)
+
+    return (intersection / (union + 1e-8)).mean()
+
+def IoU_score(model, classifier, loader, device, ttype):
+    with torch.no_grad():
+        average_recall = 0
+        for y, img in tqdm(loader):
+            out = infer_model(classifier, device, img, loader, model, ttype).squeeze()
+            y = y.squeeze().to(device)
+            recall = distribution_iou(out, y)
+            average_recall += recall.item()
+        average_recall /= len(loader)
+    print("Average IoU score:", average_recall)

vla/benchmarks/class_distribution_using_segmentation/train.py

@@ -9,6 +9,7 @@
 from model import SimpleClassifier
 from features import FeatureStore
 from precompute import precompute_features
+from metrics import score_loss, IoU_score, symmetric_topk_recall_score
 
 ttype = torch.float32
 
@@ -112,80 +113,6 @@ def train_classifier(
 
     return classifier
 
-def score_loss(model, classifier, loader, device):
-    criterion = nn.KLDivLoss(reduction='batchmean')
-    with torch.no_grad():
-        sum_loss = 0
-        for y, img in tqdm(loader):
-            if loader.dataset.feature_store is None:
-                feats = model.encode(img.to(device))
-            else:
-                feats = img.to(device)
-            feats = feats.to(ttype)
-            out = classifier(feats)
-            log_probs = F.log_softmax(out, dim=1)
-            loss = criterion(log_probs, y.to(device))
-            sum_loss += loss.item()
-        sum_loss /= len(loader)
-
-    print("Average loss:", sum_loss)
-
-def symmetric_topk_mass_recall(logits, targets, k=10):
-
-    probs = F.softmax(logits, dim=-1)
-
-    topk_pred_idx = probs.topk(k, dim=-1).indices
-    target_mass_in_pred_topk = torch.gather(targets, dim=-1, index=topk_pred_idx)
-    recall_p_to_q = target_mass_in_pred_topk.sum(dim=-1)
-
-    topk_target_idx = targets.topk(k, dim=-1).indices
-    pred_mass_in_target_topk = torch.gather(probs, dim=-1, index=topk_target_idx)
-    recall_q_to_p = pred_mass_in_target_topk.sum(dim=-1)
-
-    symmetric_recall = 0.5 * (recall_p_to_q + recall_q_to_p)
-
-    return symmetric_recall.mean()
-
-def symmetric_topk_recall_score(model, classifier, loader, device):
-    with torch.no_grad():
-        average_recall = 0
-        for y, img in tqdm(loader):
-            if loader.dataset.feature_store is None:
-                feats = model.encode(img.to(device))
-            else:
-                feats = img.to(device)
-            feats = feats.to(ttype)
-            out = classifier(feats).squeeze()
-            y = y.squeeze().to(device)
-            recall = symmetric_topk_mass_recall(out, y, k=17)
-            average_recall += recall.item()
-        average_recall /= len(loader)
-    print("Average symmetric top-k recall:", average_recall)
-
-def distribution_iou(logits, targets):
-    probs = F.softmax(logits, dim=-1)
-
-    intersection = torch.minimum(probs, targets).sum(dim=-1)
-    union = torch.maximum(probs, targets).sum(dim=-1)
-
-    return (intersection / (union + 1e-8)).mean()
-
-def IoU_score(model, classifier, loader, device):
-    with torch.no_grad():
-        average_recall = 0
-        for y, img in tqdm(loader):
-            if loader.dataset.feature_store is None:
-                feats = model.encode(img.to(device))
-            else:
-                feats = img.to(device)
-            feats = feats.to(ttype)
-            out = classifier(feats).squeeze()
-            y = y.squeeze().to(device)
-            recall = distribution_iou(out, y)
-            average_recall += recall.item()
-        average_recall /= len(loader)
-    print("Average IoU score:", average_recall)
-
 def benchmark(model, preprocessor, train_json="train_los_dataset.json", test_json="test_los_dataset.json",
               use_precomputed_features=True, random_seed=None,
               generalization_set_folder="", config_path="example_config.json"):
@@ -253,11 +180,11 @@ def benchmark(model, preprocessor, train_json="train_los_dataset.json", test_jso
 
     classifier = train_classifier(model, classifier, train_loader, test_loader, device)
     print("Computing score on test set:\n")
-    symmetric_topk_recall_score(model, classifier, score_loader, device)
-    IoU_score(model, classifier, score_loader, device)
-    score_loss(model, classifier, score_loader, device)
+    symmetric_topk_recall_score(model, classifier, score_loader, device, ttype)
+    IoU_score(model, classifier, score_loader, device, ttype)
+    score_loss(model, classifier, score_loader, device, ttype)
     if generalization_dataset_manager is not None:
         print("Computing score on generalization set:\n")
-        symmetric_topk_recall_score(model, classifier, generalization_loader, device)
-        IoU_score(model, classifier, generalization_loader, device)
-        score_loss(model, classifier, generalization_loader, device)
+        symmetric_topk_recall_score(model, classifier, generalization_loader, device, ttype)
+        IoU_score(model, classifier, generalization_loader, device, ttype)
+        score_loss(model, classifier, generalization_loader, device, ttype)