Update graph and adversarial regularization metrics to include scaled values for consistency with their respective loss term contributions.

PiperOrigin-RevId: 322626309

Author: arjung

neural_structured_learning/keras/adversarial_regularization.py

@@ -681,8 +681,10 @@ def call(self, inputs, **kwargs):
         labeled_loss=labeled_loss,
         gradient_tape=tape,
         model_kwargs=kwargs)
-    self.add_loss(self.adv_config.multiplier * adv_loss)
-    self.add_metric(adv_loss, name='adversarial_loss', aggregation='mean')
+    scaled_adv_loss = self.adv_config.multiplier * adv_loss
+    self.add_loss(scaled_adv_loss)
+    self.add_metric(
+        scaled_adv_loss, name='scaled_adversarial_loss', aggregation='mean')
     return outputs
 
   def save(self, *args, **kwargs):

neural_structured_learning/keras/graph_regularization.py

@@ -132,8 +132,10 @@ def call(self, inputs, training=False, **kwargs):
     else:
       graph_loss = tf.constant(0, dtype=tf.float32)
 
+    scaled_graph_loss = self.graph_reg_config.multiplier * graph_loss
     # Note that add_metric() cannot be invoked in a control flow branch.
-    self.add_metric(graph_loss, name='graph_loss', aggregation='mean')
-    self.add_loss(self.graph_reg_config.multiplier * graph_loss)
+    self.add_metric(
+        scaled_graph_loss, name='scaled_graph_loss', aggregation='mean')
+    self.add_loss(scaled_graph_loss)
 
     return base_output