doc string update

official/vision/beta/projects/yolo/losses/yolo_loss.py

@@ -22,9 +22,11 @@ from official.vision.beta.projects.yolo.ops import (loss_utils,
 
 
 class YoloLossBase(object, metaclass=abc.ABCMeta):
-  """Parameters for the YOLO loss functions used at each detection  
-  generator. This base class implements the base functionality required to 
-  implement a Yolo Loss function"""
+  """Parameters for the YOLO loss functions used at each detection generator. 
+  
+  This base class implements the base functionality required to implement a Yolo 
+  Loss function.
+  """
 
   def __init__(self,
                classes,
@@ -110,8 +112,7 @@ class YoloLossBase(object, metaclass=abc.ABCMeta):
     self._build_per_path_attributes()
 
   def box_loss(self, true_box, pred_box, darknet=False):
-    """Calls the iou functions and uses it to compute the loss this op is 
-    the same regardless of Yolo Loss version"""
+    """Call iou function and use it to compute the loss for the box maps."""
     if self._loss_type == "giou":
       iou, liou = box_ops.compute_giou(true_box, pred_box)
     elif self._loss_type == "ciou":
@@ -129,8 +130,7 @@ class YoloLossBase(object, metaclass=abc.ABCMeta):
                                true_conf,
                                smoothed,
                                scale=None):
-    """Completes a search of all predictions against all the ground truths to 
-    dynamically associate ground truths with predictions."""
+    """Search of all groundtruths to associate groundtruths to predictions."""
 
     # Search all predictions against ground truths to find mathcing boxes for
     # each pixel.
@@ -159,8 +159,7 @@ class YoloLossBase(object, metaclass=abc.ABCMeta):
     return true_conf, obj_mask
 
   def __call__(self, true_counts, inds, y_true, boxes, classes, y_pred):
-    """Call function to compute the loss and return the total loss as 
-    well as the loss for each detection mask on a given FPN level. 
+    """Call function to compute the loss and a set of metrics per FPN level.
     
     Args: 
       true_counts: `Tensor` of shape [batchsize, height, width, num_anchors] 
@@ -208,8 +207,7 @@ class YoloLossBase(object, metaclass=abc.ABCMeta):
 
   @abc.abstractmethod
   def _build_per_path_attributes(self):
-    """Additional initialization required specifically for each unique YOLO 
-    loss version"""
+    """Additional initialization required for each YOLO loss version"""
     ...
 
   @abc.abstractmethod
@@ -218,9 +216,20 @@ class YoloLossBase(object, metaclass=abc.ABCMeta):
     ...
 
   def post_path_aggregation(self, loss, ground_truths, predictions):
-    """This method allows for post processing of a loss value after the loss
-    has been aggregateda across all the FPN levels. The default behavior is to 
-    pass the loss through with no alterations."""
+    """This method allows for post processing of a loss value. 
+    
+    After the loss has been aggregated across all the FPN levels some post
+    proceessing may need to occur to poroperly scale the loss. The default 
+    behavior is to pass the loss through with no alterations.
+    
+    Args: 
+      loss: `tf.float` scalar for the actual loss.
+      ground_truths: `Dict` holding all the ground truth tensors.
+      predictions: `Dict` holding all the predicted values.
+
+    Returns: 
+      loss: `tf.float` scalar for the scaled loss.
+    """
     return loss
 
   @abc.abstractmethod
@@ -231,11 +240,12 @@ class YoloLossBase(object, metaclass=abc.ABCMeta):
 
 @tf.custom_gradient
 def grad_sigmoid(values):
-  """This function scales the gradient as if a signmoid was applied to the 
-  model output. This is used in the Darknet Loss when the choosen box type 
-  is the scaled coordinate type. This function is used to match the propagated 
-  gradient to match that of the Darkent Yolov4 model. This is an Identity 
-  operation that allows us to add some estra steps to the back propagation. 
+  """This function scales the gradient as if a signmoid was applied.
+  
+  This is used in the Darknet Loss when the choosen box type is the scaled 
+  coordinate type. This function is used to match the propagated gradient to 
+  match that of the Darkent Yolov4 model. This is an Identity operation that 
+  allows us to add some extra steps to the back propagation. 
   """
   def delta(dy):
     t = tf.math.sigmoid(values)
@@ -396,7 +406,7 @@ class DarknetLoss(YoloLossBase):
 
 
 class ScaledLoss(YoloLossBase):
-  """This class implements the full logic for the scaled Yolo models. """
+  """This class implements the full logic for the scaled Yolo models."""
 
   def _build_per_path_attributes(self):
     """Paramterization of pair wise search and grid generators.
@@ -520,11 +530,22 @@ class ScaledLoss(YoloLossBase):
             ind_mask, grid_mask)
 
   def post_path_aggregation(self, loss, ground_truths, predictions):
-    """By default the model will have about 3 FPN levels {3, 4, 5}, on 
+    """This method allows for post processing of a loss value. 
+    
+    By default the model will have about 3 FPN levels {3, 4, 5}, on 
     larger model that have more like 4 or 5 FPN levels the loss needs to 
     be scaled such that the total update is scaled to the same effective 
     magintude as the model with 3 FPN levels. This helps to prevent gradient 
-    explosions."""
+    explosions.
+    
+    Args: 
+      loss: `tf.float` scalar for the actual loss.
+      ground_truths: `Dict` holding all the ground truth tensors.
+      predictions: `Dict` holding all the predicted values.
+
+    Returns: 
+      loss: `tf.float` scalar for the scaled loss.
+    """
     scale = tf.stop_gradient(3 / len(list(predictions.keys())))
     return loss * scale
 

official/vision/beta/projects/yolo/ops/loss_utils.py

@@ -19,9 +19,11 @@ from official.vision.beta.projects.yolo.ops import (box_ops, math_ops)
 
 @tf.custom_gradient
 def sigmoid_bce(y, x_prime, label_smoothing):
-  """Applies the Sigmoid Cross Entropy Loss Using the same derivative as that 
-  found in the Darknet C library. The derivative of this method is not the same 
-  as the standard binary cross entropy with logits function.
+  """Applies the Sigmoid Cross Entropy Loss.
+  
+  Implements the same derivative as that found in the Darknet C library. 
+  The derivative of this method is not the same as the standard binary cross 
+  entropy with logits function.
  
   The BCE with logits function equation is as follows: 
     x = 1 / (1 + exp(-x_prime))
@@ -67,11 +69,12 @@ def sigmoid_bce(y, x_prime, label_smoothing):
 
 
 def apply_mask(mask, x, value=0):
-  """This function is used for gradient masking. The YOLO loss function makes 
-  extensive use of dynamically shaped tensors. To allow this use case on the 
-  TPU while preserving the gradient correctly for back propagation we use this 
-  masking function to use a tf.where operation to hard set masked location to 
-  have a gradient and a value of zero. 
+  """This function is used for gradient masking. 
+  
+  The YOLO loss function makes extensive use of dynamically shaped tensors. 
+  To allow this use case on the TPU while preserving the gradient correctly 
+  for back propagation we use this masking function to use a tf.where operation 
+  to hard set masked location to have a gradient and a value of zero. 
 
   Args: 
     mask: A `Tensor` with the same shape as x used to select values of 
@@ -87,9 +90,11 @@ def apply_mask(mask, x, value=0):
 
 
 def build_grid(indexes, truths, preds, ind_mask, update=False, grid=None):
-  """This function is used to broadcast all the indexes to the correct
-  ground truth mask, used for iou detection map in the scaled loss and
-  the classification mask in the darknet loss.
+  """This function is used to broadcast elements into the output shape. 
+
+  This function is used to broadcasts a list of truths into the correct index 
+  in the output shape. This is used for the ground truth map construction in 
+  the scaled loss and the classification map in the darknet loss.
   
   Args:
     indexes: A `Tensor` for the indexes
@@ -145,11 +150,10 @@ def build_grid(indexes, truths, preds, ind_mask, update=False, grid=None):
 
 
 class GridGenerator:
-  """Grid generator that generates anchor grids that will be used 
-  in to decode the predicted boxes."""
+  """Grid generator that generates anchor grids for box decoding."""
 
   def __init__(self, anchors, masks=None, scale_anchors=None):
-    """Initialize Grid Generator
+    """Initialize Grid Generator.
  
     Args:
       anchors: A `List[List[int]]` for the anchor boxes that are used in the 
@@ -173,8 +177,7 @@ class GridGenerator:
     return
 
   def _build_grid_points(self, lwidth, lheight, anchors, dtype):
-    """Generate a grid that is used to detemine the relative centers 
-    of the bounding boxs. """
+    """Generate a grid of fixed grid edges for box center decoding."""
     with tf.name_scope('center_grid'):
       y = tf.range(0, lheight)
       x = tf.range(0, lwidth)
@@ -189,7 +192,7 @@ class GridGenerator:
     return x_y
 
   def _build_anchor_grid(self, anchors, dtype):
-    """Get the transformed anchor boxes for each dimention. """
+    """Get the transformed anchor boxes for each dimention."""
     with tf.name_scope('anchor_grid'):
       num = tf.shape(anchors)[0]
       anchors = tf.cast(anchors, dtype=dtype)
@@ -217,10 +220,10 @@ class GridGenerator:
 
 TILE_SIZE = 50
 class PairWiseSearch:
-  """This method applies a pairwise search between the ground truth 
-  and the labels. The goal is to indicate the locations where the 
-  predictions overlap with ground truth for dynamic ground 
-  truth constructions."""
+  """Apply a pairwise search between the ground truth and the labels. 
+  
+  The goal is to indicate the locations where the predictions overlap with 
+  ground truth for dynamic ground truth associations."""
 
   def __init__(self,
                iou_type='iou',
@@ -554,9 +557,10 @@ def get_predicted_box(width,
                       darknet=False,
                       box_type="original",
                       max_delta=np.inf):
-  """Decodes the predicted boxes from the model format to a usable 
-  [x, y, w, h] format for use in the loss function as well as for use 
-  within the detection generator.   
+  """Decodes the predicted boxes from the model format to a usable format. 
+
+  This function decodes the model outputs into the [x, y, w, h] format for 
+  use in the loss function as well as for use within the detection generator.   
  
   Args: 
     width: A `float` scalar indicating the width of the prediction layer.