Merge pull request #15 from ranamihir/improvements

Metric improvements

Author: ranamihir

pytorch_common/config.py

@@ -105,18 +105,8 @@ def set_pytorch_config(config: _Config) -> None:
     # Fix seed
     set_seed(config.seed)
 
-    # Check for model and classification type
-    assert (
-        config.model_type == "classification" and config.classification_type in ["binary", "multiclass", "multilabel"]
-    ) or (config.model_type == "regression" and not hasattr(config, "classification_type"))
-    if config.model_type == "regression":
-        config.classification_type = None
-
-    # TODO: Remove this after extending FocalLoss
-    if config.model_type == "classification" and config.loss_criterion == "focal-loss":
-        assert (
-            config.classification_type == "binary"
-        ), "FocalLoss is currently only supported for binary classification."
+    # Check miscellaneous configurations
+    check_and_set_misc_config(config)
 
     # Ensure GPU availability as some models are prohibitively slow on CPU
     if config.assert_gpu:
@@ -186,7 +176,6 @@ def _check_loss_and_eval_criteria(config: _Config) -> None:
     assert config.get("eval_criteria") and isinstance(config.eval_criteria, list)
 
     loss_criteria = CLASSIFICATION_LOSS_CRITERIA if config.model_type == "classification" else REGRESSION_LOSS_CRITERIA
-
     assert config.loss_criterion in loss_criteria, (
         f"Loss criterion ('{config.loss_criterion}') "
         f"for `model_type=='classification' must be one"
@@ -306,3 +295,29 @@ def set_mode_batch_size(mode: str, batch_size_per_gpu: int) -> None:
 
         # Set per-GPU and total batch size for mode
         set_mode_batch_size(mode, batch_size_to_set)
+
+
+def check_and_set_misc_config(config: _Config) -> None:
+    """
+    Check all miscellaneous configurations, e.g.:
+      - model_type
+      - classification_type
+    """
+    # Check for model and classification type
+    assert (
+        config.model_type == "classification" and config.classification_type in ["binary", "multiclass", "multilabel"]
+    ) or (config.model_type == "regression" and not hasattr(config, "classification_type"))
+
+    # Set classification_type to None if regression
+    if config.model_type == "regression":
+        config.classification_type = None
+
+    # TODO: Remove this after extending FocalLoss
+    if config.model_type == "classification" and config.loss_criterion == "focal-loss":
+        assert (
+            config.classification_type == "binary"
+        ), "FocalLoss is currently only supported for binary classification."
+
+    # Used for dataloader sampling
+    config.num_batches = config.get("num_batches", None)
+    config.percentage = config.get("percentage", None)

pytorch_common/datasets.py

@@ -1,5 +1,12 @@
+import numpy as np
+from torch.utils.data import DataLoader, Dataset
+
 from .additional_configs import BaseDatasetConfig
 from .datasets_dl import BasePyTorchDataset, DummyMultiClassDataset, DummyMultiLabelDataset, DummyRegressionDataset
+from .types import Optional, _Config
+from .utils import setup_logging
+
+logger = setup_logging(__name__)
 
 
 def create_dataset(dataset_name: str, config: BaseDatasetConfig) -> BasePyTorchDataset:
@@ -16,3 +23,46 @@ def create_dataset(dataset_name: str, config: BaseDatasetConfig) -> BasePyTorchD
     else:
         raise RuntimeError(f"Unknown dataset name {dataset_name}.")
     return dataset
+
+
+def create_dataloader(dataset: Dataset, config: _Config, is_train: Optional[bool] = True) -> DataLoader:
+    """
+    Create a dataloader wrapped
+    around the given dataset.
+
+    Option to sample a subset of the data:
+    During development, you can just set
+    `num_batches` or `percentage` to a small
+    number to run quickly on a sample dataset.
+    """
+    if is_train:
+        shuffle = True
+        batch_size = config.train_batch_size
+    else:
+        shuffle = False
+        batch_size = config.eval_batch_size
+
+    num_batches, percentage = config.num_batches, config.percentage
+
+    if num_batches and percentage:
+        raise ValueError(
+            f"At most one of `num_batches` ({num_batches}) or `percentage` ({percentage}) may be specified."
+        )
+
+    elif num_batches or percentage:
+        n = len(dataset)
+        if num_batches:
+            assert num_batches <= np.ceil(n / batch_size).astype(int)
+            logger.info(f"Sampling {num_batches} batches from whole dataloader.")
+            sampled_indices = np.random.choice(range(n), size=min(num_batches * batch_size, n))
+        else:
+            assert percentage <= 100.0
+            logger.info(f"Sampling {percentage}% of whole dataset.")
+            sampled_indices = np.random.choice(range(n), size=int(percentage * n))
+
+        dataset.data = dataset.data.iloc[sampled_indices]
+        logger.info("Done.")
+
+    dataloader = DataLoader(dataset, batch_size=batch_size, shuffle=shuffle, num_workers=8, pin_memory=True)
+
+    return dataloader

pytorch_common/decorators.py

@@ -43,9 +43,9 @@ def timing_with_param(*parameter_names) -> Callable:
     def decorator(func: Callable) -> Callable:
         @wraps(func)
         def wrapper(*args, **kwargs):
-            start = time.time()
+            start = time.perf_counter()
             result = func(*args, **kwargs)
-            end = time.time()
+            end = time.perf_counter()
             elapsed = end - start
             elapsed_human = human_time_interval(elapsed)
 
@@ -54,7 +54,7 @@ def wrapper(*args, **kwargs):
             logged_param_str = f" {logged_param}" if logged_param else ""
 
             module_name, function_name = func.__module__, func.__qualname__
-            PRINT_FUNC(f"Function '{module_name}.{function_name}{logged_param_str}' took {elapsed_human}")
+            PRINT_FUNC(f"Function '{module_name}.{function_name}{logged_param_str}' took {elapsed_human}.")
             return result
 
         return wrapper

pytorch_common/metric_utils.py

@@ -58,6 +58,7 @@ def correct_argument_order(func: Callable) -> Callable:
     return corrected_func
 
 
+@torch.no_grad()
 def get_mse_loss(y_predicted: torch.Tensor, y_true: torch.Tensor, **kwargs) -> float:
     """
     Compute MSE loss.
@@ -75,14 +76,56 @@ def auc_score(y_predicted: torch.Tensor, y_true: torch.Tensor, **kwargs) -> floa
     return auc(fpr, tpr)
 
 
+@torch.no_grad()
+def top_k_accuracy_scores(
+    eval_metrics: List[object], results: _StringDict, y_predicted: torch.Tensor, y_true: torch.Tensor, **kwargs
+) -> None:
+    """
+    A batch implementation of `top_k_accuracy_score()`.
+
+    It takes all top_k_accuracy-based metrics as input,
+    and computes the respective metrics efficiently.
+    If done separately, the top-k indices would need
+    to be computed separately for each k, while here
+    it happens only once.
+    """
+    assert len(y_predicted) == len(y_true)
+
+    ks = []
+    for eval_metric in eval_metrics:
+        k = int(match(eval_metric.criterion, TOP_K_ACCURACY_REGEX))  # Extract k
+        ks.append(k)
+
+    max_k = max(ks)
+    _, top_indices = torch.topk(y_predicted, max_k, dim=1)  # Compute the top `max_k` predicted classes
+    top_indices = top_indices.t()  # Transpose for mathematical convenience
+    correct_max_k = top_indices.eq(
+        y_true.long().view(1, -1).expand_as(top_indices)
+    )  # Get correct predictions in top max_k
+
+    # Compute top-k accuracy for all k's
+    for i, k in enumerate(ks):
+        correct_k = correct_max_k[:k].reshape(-1).float().sum(dim=0, keepdim=True)  # Get correct predictions in top k
+        top_k_accuracy = correct_k / len(y_true)  # Divide by batch size (because of transpose earlier)
+        results[eval_metrics[i].criterion] = top_k_accuracy.item()
+
+
+@torch.no_grad()
 def top_k_accuracy_score(y_predicted: torch.Tensor, y_true: torch.Tensor, **kwargs) -> float:
     """
     Compute the top-k accuracy score
     in a multi-class setting.
 
     Conversion to numpy is expensive in this
     case. Stick to using PyTorch tensors.
+
+    Note: This function is not recommended if you have
+          more than one k that this is to be computed
+          for. Please use the much more efficient
+          `top_k_accuracy_scores()` in that case.
     """
+    assert len(y_predicted) == len(y_true)
+
     k = int(match(kwargs["criterion"], TOP_K_ACCURACY_REGEX))  # Extract k
     _, topk_indices = torch.topk(y_predicted, k, dim=1)  # Compute the top-k predicted classes
     correct_examples = torch.eq(y_true[..., None, ...].long(), topk_indices).any(dim=1)
@@ -191,7 +234,7 @@ def __repr__(self):
     },
     "auc": {"preprocess_fn": prob_class1, "eval_fn": auc_score, "model_type": "classification"},
     "top_k_accuracy": {
-        "eval_fn": top_k_accuracy_score,
+        "eval_fn": top_k_accuracy_score,  # Not actually used (in favor of `top_k_accuracy_scores()` for efficiecy)
         "regex": TOP_K_ACCURACY_REGEX,
         "model_type": "classification",
     },

pytorch_common/metrics.py

@@ -5,7 +5,15 @@
 import torch
 import torch.nn as nn
 
-from .metric_utils import EVAL_METRIC_FUNCTIONS, LOSS_CRITERIA, PREPROCESSING_FUNCTIONS, FocalLoss, canonicalize, match
+from .metric_utils import (
+    EVAL_METRIC_FUNCTIONS,
+    LOSS_CRITERIA,
+    PREPROCESSING_FUNCTIONS,
+    FocalLoss,
+    canonicalize,
+    match,
+    top_k_accuracy_scores,
+)
 from .types import *
 
 
@@ -232,12 +240,26 @@ def compute_per_class(self, y_predicted: torch.Tensor, y_true: torch.Tensor) ->
         """
         results = {}
         for preprocess_fn, supported_metrics in PREPROCESSING_FUNCTIONS.items():
-            metrics_to_compute = [metric_name for metric_name in supported_metrics if metric_name in self.criteria]
-            if metrics_to_compute:
+            metrics_in_group = [metric_name for metric_name in supported_metrics if metric_name in self.criteria]
+            if metrics_in_group:
                 preprocessed_input = preprocess_fn(y_predicted, y_true)  # Share preprocessing
-                for metric_name in metrics_to_compute:
+                for metric_name in metrics_in_group:
+                    # Get eval metrics to be computed in this group
                     eval_metrics = [criterion for criterion in self.criteria if criterion.name == metric_name]
+
+                    # Separate out all top_k_accuracy-based metrics
+                    top_k_accuracy_metrics, other_metrics = [], []
                     for eval_metric in eval_metrics:
+                        top_k_accuracy_metrics.append(
+                            eval_metric
+                        ) if eval_metric.name == "top_k_accuracy" else other_metrics.append(eval_metric)
+
+                    # Compute all top_k_accuracy-based metrics together (for efficiency)
+                    for eval_metric in top_k_accuracy_metrics:
+                        top_k_accuracy_scores(top_k_accuracy_metrics, results, *preprocessed_input)
+
+                    # Compute all other metrics as normal
+                    for eval_metric in other_metrics:
                         results[eval_metric.criterion] = eval_metric(*preprocessed_input)
         return results