Add graph and adversarial loss values scalars to the summary.

This will allow visualization of the graph and adversrial losses in TensorBoard.

PiperOrigin-RevId: 321660205

Author: arjung

neural_structured_learning/estimator/adversarial_regularization.py

@@ -121,12 +121,18 @@ def adv_model_fn(features, labels, mode, params=None, config=None):
 
       # Runs the base model again to compute loss on adv_neighbor.
       adv_spec = base_fn(adv_neighbor, labels)
+      scaled_adversarial_loss = adv_config.multiplier * adv_spec.loss
+      tf.compat.v1.summary.scalar('loss/scaled_adversarial_loss',
+                                  scaled_adversarial_loss)
 
-      final_loss = original_spec.loss + adv_config.multiplier * adv_spec.loss
+      supervised_loss = original_spec.loss
+      tf.compat.v1.summary.scalar('loss/supervised_loss', supervised_loss)
+
+      final_loss = supervised_loss + scaled_adversarial_loss
 
       if not optimizer_fn:
         # Default to the Adagrad optimizer, the same as canned DNNEstimator.
-        optimizer = tf.train.AdagradOptimizer(learning_rate=0.05)
+        optimizer = tf.compat.v1.train.AdagradOptimizer(learning_rate=0.05)
       else:
         optimizer = optimizer_fn()
 

neural_structured_learning/estimator/graph_regularization.py

@@ -144,11 +144,17 @@ def graph_reg_model_fn(features, labels, mode, params=None, config=None):
           nbr_embeddings,
           weights=nbr_weights,
           distance_config=graph_reg_config.distance_config)
-      total_loss = base_spec.loss + graph_reg_config.multiplier * graph_loss
+      scaled_graph_loss = graph_reg_config.multiplier * graph_loss
+      tf.compat.v1.summary.scalar('loss/scaled_graph_loss', scaled_graph_loss)
+
+      supervised_loss = base_spec.loss
+      tf.compat.v1.summary.scalar('loss/supervised_loss', supervised_loss)
+
+      total_loss = supervised_loss + scaled_graph_loss
 
       if not optimizer_fn:
         # Default to Adagrad optimizer, the same as the canned DNNEstimator.
-        optimizer = tf.train.AdagradOptimizer(learning_rate=0.05)
+        optimizer = tf.compat.v1.train.AdagradOptimizer(learning_rate=0.05)
       else:
         optimizer = optimizer_fn()
       train_op = optimizer.minimize(