Added more documentation.

Author: otiliastr

neural_structured_learning/research/gam/trainer/adversarial.py

@@ -7,6 +7,7 @@
 
 
 def kl_divergence_with_logit(q_logit, p_logit):
+    """Computes KL-divergence between to sets of logits."""
     q = tf.nn.softmax(q_logit)
     qlogq = -tf.nn.softmax_cross_entropy_with_logits_v2(
         labels=q, logits=q_logit)
@@ -28,13 +29,26 @@ def get_normalizing_constant(d):
 
 
 def get_loss_vat(inputs, predictions, is_train, model, predictions_var_scope):
+    """Computes the virtual adversarial loss for the provided inputs.
+
+    Args:
+        inputs: A batch of input features, where the batch is the first
+            dimension.
+        predictions: The logits predicted by a model on the provided inputs.
+        is_train: A boolean placeholder specifying if this is a training or
+            testing setting.
+        model: The model that generated the logits.
+        predictions_var_scope: Variable scope for obtaining the predictions.
+    Returns:
+        A float value representing the virtual adversarial loss.
+    """
     r_vadv = generate_virtual_adversarial_perturbation(
         inputs, predictions, model, predictions_var_scope, is_train=is_train)
     predictions = tf.stop_gradient(predictions)
     logit_p = predictions
     new_inputs = tf.add(inputs, r_vadv)
     with tf.variable_scope(
-        predictions_var_scope,  auxiliary_name_scope=False, reuse=True):
+            predictions_var_scope,  auxiliary_name_scope=False, reuse=True):
         encoding_m, _, _ = model.get_encoding_and_params(
             inputs=new_inputs,
             is_train=is_train,
@@ -98,5 +112,5 @@ def logsoftmax(x):
 def entropy_y_x(logit):
     """Entropy term to add to VATENT."""
     p = tf.nn.softmax(logit)
-    return tf.reduce_mean(tf.nn.softmax_cross_entropy_with_logits_v2(
-        labels=p, logits=logit))
+    return tf.reduce_mean(
+        tf.nn.softmax_cross_entropy_with_logits_v2(labels=p, logits=logit))