yolo input test

official/vision/beta/projects/yolo/configs/experiments/yolov4-csp/inference/640.yaml

@@ -27,9 +27,11 @@ task:
         'all': 2.0
       max_boxes: 300
       nms_type: greedy
-      iou_thresh: 0.25
-      nms_thresh: 0.45
-      pre_nms_points: 500
+      # iou_thresh: 0.25
+      # nms_thresh: 0.45
+      iou_thresh: 0.001
+      nms_thresh: 0.60
+      pre_nms_points: 5000
     loss:
       use_scaled_loss: true
       update_on_repeat: true
@@ -70,5 +72,5 @@ task:
       aug_rand_translate: 0.1
       area_thresh: 0.1
   validation_data:
-    global_batch_size: 1
+    global_batch_size: 8
     input_path: '/media/vbanna/DATA_SHARE/CV/datasets/COCO_raw/records/val*'
\ No newline at end of file

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

+from official.vision.beta.projects.yolo.common import registry_imports  # pylint: disable=unused-import
+from official.vision.beta.projects.yolo.tasks import image_classification as imc
+from official.vision.beta.projects.yolo.configs import darknet_classification as dcfg
+
+import os
+import tensorflow as tf
+from official.core import train_utils
+from official.core import task_factory
+from absl.testing import parameterized
+
+PATH_TO_COCO = '/media/vbanna/DATA_SHARE/CV/datasets/COCO_raw/records/'
+
+def test_yolo_input_task(scaled_pipeline = False, batch_size = 1):
+  if not scaled_pipeline:
+    experiment = "yolo_darknet"
+    config_path = [
+      "official/vision/beta/projects/yolo/configs/experiments/yolov4/tpu/512.yaml"]
+  else:
+    experiment = "scaled_yolo"
+    config_path = [
+      "official/vision/beta/projects/yolo/configs/experiments/yolov4-csp/tpu/640.yaml"]
+
+  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)
+
+  config.train_data.global_batch_size = batch_size
+  config.validation_data.global_batch_size = 1
+  config.train_data.dtype = 'float32'
+  config.validation_data.dtype = 'float32'
+  config.validation_data.shuffle_buffer_size = 1
+  config.train_data.shuffle_buffer_size = 1
+  config.train_data.input_path = os.path.join(PATH_TO_COCO, 'train*')
+  config.validation_data.input_path = os.path.join(PATH_TO_COCO, 'val*')
+
+  with tf.device('/CPU:0'):
+    train_data = task.build_inputs(config.train_data)
+    test_data = task.build_inputs(config.validation_data)
+  return train_data, test_data, config
+
+def test_yolo_pipeline_visually(is_training=True, num=30):
+  # visualize the datapipeline
+  import matplotlib.pyplot as plt
+  dataset, testing, _ = test_yolo_input_task()
+
+  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]])
+
+    gt = label['true_conf']
+
+    obj3 = tf.clip_by_value(gt['3'][..., 0], 0.0, 1.0)
+    obj4 = tf.clip_by_value(gt['4'][..., 0], 0.0, 1.0)
+    obj5 = tf.clip_by_value(gt['5'][..., 0], 0.0, 1.0)
+
+    for shind in range(1):
+      fig, axe = plt.subplots(1, 4)
+
+      image = image[shind]
+
+      axe[0].imshow(image)
+      axe[1].imshow(obj3[shind].numpy())
+      axe[2].imshow(obj4[shind].numpy())
+      axe[3].imshow(obj5[shind].numpy())
+
+      fig.set_size_inches(18.5, 6.5, forward=True)
+      plt.tight_layout()
+      plt.show()
+
+class YoloDetectionInputTest(tf.test.TestCase, parameterized.TestCase):
+
+  @parameterized.named_parameters(('scaled', True), ('darknet', False))
+  def test_yolo_input(self, scaled_pipeline):
+    # builds a pipline forom the config and tests the datapipline shapes
+    dataset, _, params = test_yolo_input_task(
+        scaled_pipeline=scaled_pipeline, 
+        batch_size=1)
+
+    dataset = dataset.take(1)
+
+    for image, label in dataset:
+      self.assertAllEqual(image.shape, ([1] + params.model.input_size))
+      self.assertTrue(
+          tf.reduce_all(tf.math.logical_and(image >= 0, image <= 1)))
+
+
+if __name__ == '__main__':
+  tf.test.main()
+  # test_yolo_pipeline_visually(is_training=True, num=20)

official/vision/beta/projects/yolo/modeling/yolo_model.py

@@ -113,6 +113,7 @@ class Yolo(tf.keras.Model):
     return weights, bias, other
   
   def fuse(self):
+    """Fuses all Convolution and Batchnorm layers to get better latency."""
     print("Fusing Conv Batch Norm Layers.")
     if not self._fused:
       self._fused = True

official/vision/beta/projects/yolo/optimization/sgd_torch.py

@@ -30,27 +30,35 @@ def _var_key(var):
 
 
 class SGDTorch(tf.keras.optimizers.Optimizer):
-  """Optimizer that computes an exponential moving average of the variables.
-
-  Empirically it has been found that using the moving average of the trained
-  parameters of a deep network is better than using its trained parameters
-  directly. This optimizer allows you to compute this moving average and swap
-  the variables at save time so that any code outside of the training loop
-  will use by default the average values instead of the original ones.
+  """Optimizer that simulates the SGD module used in pytorch. 
+  
+  
+  For details on the differences between the original SGD implemention and the 
+  one in pytorch: https://pytorch.org/docs/stable/generated/torch.optim.SGD.html.
+  This optimizer also allow for the usage of a momentum warmup along side a 
+  learning rate warm up, though using this is not required. 
 
   Example of usage for training:
   ```python
-  opt = tf.keras.optimizers.SGD(learning_rate)
-  opt = ExponentialMovingAverage(opt)
+  opt = SGDTorch(learning_rate)
 
-  opt.shadow_copy(model)
-  ```
+  # Models must implement a method to iterate all model.trainable_variables
+  # and split the variables by key into the weights, biases, and others.
+  # Weight decay wil be applied to all variables in the weights group. Bias 
+  # and others are included as a way to proved alternate LR scedules to various 
+  # paramter groups. An example of this variable search can be found in 
+  # official/vision/beta/projects/yolo/modeling/yolo_model.py.
 
-  At test time, swap the shadow variables to evaluate on the averaged weights:
-  ```python
-  opt.swap_weights()
-  # Test eval the model here
-  opt.swap_weights()
+  weights, biases, other = model.get_groups()
+  opt.set_params(weights, biases, other)
+
+  # if the learning rate schedule on the biases are different. if lr is not set 
+  # the default schedule used for weights will be used on the biases. 
+  opt.set_bias_lr(<lr schedule>)
+
+  # if the learning rate schedule on the others are different. if lr is not set 
+  # the default schedule used for weights will be used on the biases. 
+  opt.set_other_lr(<lr schedule>)
   ```
   """