addressing comments

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

+runtime:
+  distribution_strategy: 'tpu'
+  mixed_precision_dtype: 'bfloat16'
+task:
+  coco91to80: True
+  smart_bias_lr: 0.1
+  reduced_logs: True
+  model:
+    base:
+      backbone:
+        type: 'darknet'
+        darknet:
+          model_id: 'altered_cspdarknet53'
+      decoder:
+        version: v4
+        type: csp
+    num_classes: 80
+    smart_bias: true
+    boxes: ['[12.0, 16.0]',
+            '[19.0, 36.0]',
+            '[40.0, 28.0]',
+            '[36.0, 75.0]', 
+            '[76.0, 55.0]', 
+            '[72.0, 146.0]', 
+            '[142.0, 110.0]', 
+            '[192.0, 243.0]', 
+            '[459.0, 401.0]']
+    input_size: [640, 640, 3]
+    boxes_per_scale: 3
+    max_level: 5
+    min_level: 3
+    norm_activation:
+      activation: mish
+      norm_epsilon: 0.0001
+      norm_momentum: 0.97
+      use_sync_bn: true
+    filter:
+      loss_type:  
+        'all': ciou
+      iou_normalizer: 
+        'all': 0.05
+      cls_normalizer: 
+        'all': 0.3
+      ignore_thresh:
+        'all': 0.0
+      obj_normalizer: 
+        '5': 0.28
+        '4': 0.70
+        '3': 2.80
+      new_cords:
+        'all': True
+      objectness_smooth: 
+        'all': 1.0
+      scale_xy:
+        'all': 2.0
+      max_boxes: 300
+      nms_type: iou
+      iou_thresh: 0.001
+      nms_thresh: 0.60
+      pre_nms_points: 5000
+      use_scaled_loss: true
+      update_on_repeat: true
+  train_data:
+    global_batch_size: 64
+    dtype: float32
+    input_path: 'gs://cam2-datasets/coco/train*'
+    is_training: true
+    shuffle_buffer_size: 10
+    drop_remainder: true
+    seed: null
+    parser:
+      max_num_instances: 300
+      letter_box: True
+      random_flip: True
+      aug_rand_saturation: 0.7
+      aug_rand_brightness: 0.4
+      aug_rand_hue: 0.015
+      aug_scale_min: 0.1
+      aug_scale_max: 1.9
+      aug_rand_translate: 0.1
+      area_thresh: 0.1
+      random_pad: False
+      use_tie_breaker: True
+      use_scale_xy: True
+      anchor_thresh: 4.0
+      best_match_only: True
+      mosaic: 
+        output_resolution: [640, 640]
+        mosaic_frequency: 1.0
+        mixup_frequency: 0.0
+        mosaic_crop_mode: 'scale'
+        aspect_ratio_mode: 'letter'
+        crop_area_mosaic: [0.1, 1.9]
+        crop_area: [0.25]
+        random_pad: False
+  validation_data:
+    global_batch_size: 64
+    dtype: float32
+    input_path: 'gs://cam2-datasets/coco/val*'
+    is_training: false
+    shuffle_buffer_size: 10000
+    drop_remainder: true
+    parser:
+      max_num_instances: 300
+      letter_box: True
+      use_tie_breaker: True
+      use_scale_xy: True
+      anchor_thresh: 4.0
+      best_match_only: True
+  weight_decay: 0.000
+  annotation_file: null
+trainer:
+  train_steps: 555000 # 160 epochs at 64 batchsize -> 500500 * 64/2
+  validation_steps: 78
+  steps_per_loop: 1850
+  summary_interval: 1850
+  validation_interval: 1850
+  checkpoint_interval: 1850
+  optimizer_config:
+    ema:
+      average_decay: 0.9999
+      trainable_weights_only: False
+      dynamic_decay: True
+    learning_rate:
+      type: cosine
+      cosine:
+        initial_learning_rate: 0.01
+        name: Cosine
+        alpha: 0.2
+        decay_steps: 555000
+    optimizer:
+      type: sgd_dymow
+      sgd_dymow:
+        momentum: 0.937
+        momentum_start: 0.8
+        nesterov: True
+        warmup_steps: 5550
+        weight_decay: 0.0005
+        sim_torch: true
+        name: SGD
+    warmup:
+      type: 'linear'
+      linear:
+        warmup_steps: 5550 #learning rate rises from 0 to 0.0013 over 1000 steps

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

@@ -38,14 +38,14 @@ def build_yolo_decoder(input_specs, model_config: yolo.Yolo, l2_regularization):
 
   if model_config.decoder.version not in yolo_model.YOLO_MODELS.keys():
     raise Exception(
-        "unsupported model version please select from {v3, v4}, \n\n \
+        "unsupported model version please select from {v3, v4}, \
         or specify a custom decoder config using YoloDecoder in you yaml")
 
   if model_config.decoder.type not in yolo_model.YOLO_MODELS[
       model_config.decoder.version].keys():
     raise Exception("unsupported model type please select from \
         {yolo_model.YOLO_MODELS[model_config.decoder.version].keys()},\
-        \n\n or specify a custom decoder config using YoloDecoder in you yaml")
+        or specify a custom decoder config using YoloDecoder in you yaml")
 
   base_model = yolo_model.YOLO_MODELS[model_config.decoder.version][
       model_config.decoder.type]
@@ -74,9 +74,26 @@ def build_yolo_decoder(input_specs, model_config: yolo.Yolo, l2_regularization):
 def build_yolo_filter(model_config: yolo.Yolo, decoder: YoloDecoder, masks,
                       xy_scales, path_scales):
   def _build(values):
-    """Used to make model cfg shorter when different FPN lvls require 
+    """Used to make model config shorter when different FPN levels require 
     different hyper paramters while others do not. The term all will 
-    blanket set the same hyper paramter for all FPN levels."""
+    blanket set the same hyper paramter for all FPN levels.
+    
+    Example: 
+      cls_normalizer: 
+        'all': 0.3
+      obj_normalizer: 
+        '5': 0.28 
+        '4': 0.70
+        '3': 2.80
+
+    Under normal use cases the config is extended by 3x because the values for 
+    levels 3, 4, and 5 need to be specified individually. Using this function 
+    allows us to specify the value "all" to indicate that all fpn levels will 
+    use the same value. 
+
+    An exmaple config can be found at: 
+    official/vision/beta/projects/yolo/configs/experiments/yolov4-csp/640.yaml
+    """
     if "all" in values and values["all"] is not None:
       for key in values:
         if key != 'all':

official/vision/beta/projects/yolo/modeling/layers/detection_generator.py

@@ -79,6 +79,8 @@ class YoloLayer(tf.keras.Model):
       nms_type: `str` for which non max suppression to use.
       objectness_smooth: `float` for how much to smooth the loss on the
         detection map.
+      kwargs**: `Dict` constining additional inputs used for the initialization 
+        of the Keras Model. 
 
     Return:
       loss: `float` for the actual loss.

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

@@ -15,10 +15,8 @@
 # Lint as: python3
 """Tests for Yolo models."""
 
-from official.vision.beta.projects.yolo.modeling.backbones.darknet import \
-    Darknet
-from official.vision.beta.projects.yolo.modeling.decoders.yolo_decoder import \
-    YoloDecoder
+from official.vision.beta.projects.yolo.modeling.backbones.darknet import Darknet
+from official.vision.beta.projects.yolo.modeling.decoders.yolo_decoder import YoloDecoder
 from official.vision.beta.projects.yolo.modeling.heads.yolo_head import YoloHead
 from official.vision.beta.projects.yolo.modeling import yolo_model