yolo task test

official/vision/beta/projects/yolo/dataloaders/yolo_input_test.py

@@ -20,7 +20,8 @@ def test_yolo_input_task(scaled_pipeline = False, batch_size = 1):
     config_path = [
       "official/vision/beta/projects/yolo/configs/experiments/yolov4-csp/tpu/640.yaml"]
 
-  config = train_utils.ParseConfigOptions(experiment=experiment, config_file=config_path)
+  config = train_utils.ParseConfigOptions(experiment=experiment, 
+                                          config_file=config_path)
   params = train_utils.parse_configuration(config)
   config = params.task
   task = task_factory.get_task(params.task)
@@ -47,7 +48,8 @@ def test_yolo_pipeline_visually(is_training=True, num=30):
   data = dataset if is_training else testing
   data = data.take(num)
   for l, (image, label) in enumerate(data):
-    image = tf.image.draw_bounding_boxes(image, label['bbox'], [[1.0, 0.0, 1.0]])
+    image = tf.image.draw_bounding_boxes(
+        image, label['bbox'], [[1.0, 0.0, 1.0]])
 
     gt = label['true_conf']
 

official/vision/beta/projects/yolo/tasks/yolo_test.py

+from official.vision.beta.projects.yolo.common import registry_imports
+
+import orbit
+from absl.testing import parameterized
+from official.core import exp_factory
+from official.modeling import optimization
+
+from official.modeling.optimization import configs
+from official.core import train_utils
+
+from official.vision.beta.projects.yolo.tasks import yolo
+
+import tensorflow as tf
+
+
+class YoloTaskTest(tf.test.TestCase, parameterized.TestCase):
+
+  @parameterized.parameters(("scaled_yolo",))
+  def test_task(self, config_name):
+    config_path = ["official/vision/beta/projects/yolo/configs/experiments/yolov4-csp/inference/640.yaml"]
+    config = exp_factory.get_exp_config(config_name)
+
+    config = train_utils.ParseConfigOptions(
+        experiment=config_name, config_file=config_path)
+    config = train_utils.parse_configuration(config)
+
+    config.trainer.optimizer_config.ema = None
+    config.task.train_data.global_batch_size = 1
+    config.task.validation_data.global_batch_size = 1
+
+    task = yolo.YoloTask(config.task)
+    model = task.build_model()
+    metrics = task.build_metrics(training=False)
+    strategy = tf.distribute.get_strategy()
+
+    train = orbit.utils.make_distributed_dataset(strategy, task.build_inputs,
+                                                   config.task.train_data)
+    test = orbit.utils.make_distributed_dataset(strategy, task.build_inputs,
+                                                   config.task.validation_data)
+    train = iter(train)
+    test = iter(test)
+    optimizer = task.create_optimizer(config.trainer.optimizer_config)
+    logs = task.train_step(next(train), model, optimizer, metrics=metrics)
+    self.assertIn("loss", logs)
+    logs = task.validation_step(next(test), model, metrics=metrics)
+    self.assertIn("loss", logs)
+
+
+if __name__ == "__main__":
+  tf.test.main()

official/vision/beta/projects/yolo/train.py

@@ -33,7 +33,7 @@ FLAGS = flags.FLAGS
 python3 -m official.vision.beta.projects.yolo.train --mode=train_and_eval --experiment=darknet_classification --model_dir=training_dir --config_file=official/vision/beta/projects/yolo/configs/experiments/darknet53_tfds.yaml
 
 
-python3.8 -m official.vision.beta.projects.yolo.train --experiment=yolo_darknet --mode train_and_eval --config_file yolo/configs/experiments/yolov4/inference/512-swin.yaml --model_dir ../checkpoints/test-swin
+python3.8 -m official.vision.beta.projects.yolo.train --experiment=scaled_yolo --mode train_and_eval --config_file yolo/configs/experiments/yolov4-csp/tpu/640.yaml --model_dir ../checkpoints/yolov4-csp-640
 '''