Adds per-class metrics.

PiperOrigin-RevId: 366540340

official/vision/beta/configs/video_classification.py

@@ -101,6 +101,11 @@ class Losses(hyperparams.Config):
   l2_weight_decay: float = 0.0
 
 
+@dataclasses.dataclass
+class Metrics(hyperparams.Config):
+  use_per_class_recall: bool = False
+
+
 @dataclasses.dataclass
 class VideoClassificationTask(cfg.TaskConfig):
   """The task config."""
@@ -109,6 +114,7 @@ class VideoClassificationTask(cfg.TaskConfig):
   validation_data: DataConfig = DataConfig(
       is_training=False, drop_remainder=False)
   losses: Losses = Losses()
+  metrics: Metrics = Metrics()
 
 
 def add_trainer(experiment: cfg.ExperimentConfig,

official/vision/beta/tasks/video_classification.py

@@ -154,6 +154,10 @@ class VideoClassificationTask(base_task.Task):
                 curve='PR',
                 multi_label=self.task_config.train_data.is_multilabel,
                 name='PR-AUC'))
+        if self.task_config.metrics.use_per_class_recall:
+          for i in range(self.task_config.train_data.num_classes):
+            metrics.append(
+                tf.keras.metrics.Recall(class_id=i, name=f'recall-{i}'))
     else:
       metrics = [
           tf.keras.metrics.SparseCategoricalAccuracy(name='accuracy'),