Skip logging the input tensors to the loss block (aws#86)

* Skip logging the input tensors to the loss block

* Add loss inputs for PT functional loss

* append mode name

* Write correct mode for each scalar in write_scalars()
* Increment global mode step number irrespective of which mode it is

Author: vandanavk

smdebug/core/hook.py

@@ -47,9 +47,10 @@
 
 
 class ScalarCache(object):
-    def __init__(self, scalar_name, scalar_val, sm_metric, write_tb, write_event):
+    def __init__(self, scalar_name, scalar_val, mode, sm_metric, write_tb, write_event):
         self.name = scalar_name
         self.value = scalar_val
+        self.mode = mode
         self.sm_metric = sm_metric
         self.write_tb = write_tb
         self.write_event = write_event
@@ -440,6 +441,10 @@ def _increment_step(self):
 
         self.step += 1
         self.mode_steps[self.mode] += 1
+
+        # Increment Global step number irrespective of what mode it is
+        if self.mode != ModeKeys.GLOBAL:
+            self.mode_steps[ModeKeys.GLOBAL] = self.step
         self._collections_to_save_for_step = None
 
     def _write_state(self):
@@ -564,12 +569,15 @@ def _write_scalars(self):
         for scalar_obj in self.scalar_cache:
             scalar_name = scalar_obj.name
             scalar_val = scalar_obj.value
+            scalar_mode = scalar_obj.mode
             sm_metric = scalar_obj.sm_metric
             write_tb = scalar_obj.write_tb
             write_event = scalar_obj.write_event
             if self.metrics_writer and sm_metric:
                 self.metrics_writer.log_metric(
-                    scalar_name, scalar_val, iteration_number=self.mode_steps[self.mode]
+                    scalar_name + "_" + scalar_mode.name,
+                    scalar_val,
+                    iteration_number=self.mode_steps[scalar_mode],
                 )
             if write_tb:
                 tb_writer = self._maybe_get_tb_writer()
@@ -596,7 +604,7 @@ def save_scalar(self, name, value, sm_metric=False):
         val = self._make_numpy_array(value)
         if val.size != 1:
             raise TypeError(f"{name} has non scalar value of type: {type(value)}")
-        scalar_obj = ScalarCache(name, val, sm_metric=True, write_tb=True, write_event=True)
+        scalar_obj = ScalarCache(name, val, self.mode, sm_metric, write_tb=True, write_event=True)
         self.scalar_cache.append(scalar_obj)
 
     def _write_raw_tensor(self, tensor_name, tensor_value, save_collections, tensor_ref=None):
@@ -657,7 +665,12 @@ def _save_for_tensor(self, tensor_name, tensor_value, check_before_write=True):
                 # Always log loss to Minerva
                 tensor_val = np.mean(np_val)
                 scalar_obj = ScalarCache(
-                    tensor_name, tensor_val, sm_metric=True, write_tb=False, write_event=False
+                    tensor_name,
+                    tensor_val,
+                    self.mode,
+                    sm_metric=True,
+                    write_tb=False,
+                    write_event=False,
                 )
                 self.scalar_cache.append(scalar_obj)
 

smdebug/mxnet/collection.py

@@ -24,7 +24,7 @@ def _register_default_collections(self):
         self.get(CollectionKeys.WEIGHTS).include("^(?!gradient).*weight")
         self.get(CollectionKeys.BIASES).include("^(?!gradient).*bias")
         self.get(CollectionKeys.GRADIENTS).include("^gradient")
-        self.get(CollectionKeys.LOSSES).include(".*loss")
+        self.get(CollectionKeys.LOSSES).include(".*loss._(?!input).*output")
 
     def create_collection(self, name):
         super().create_collection(name, cls=Collection)

smdebug/pytorch/collection.py

@@ -39,7 +39,7 @@ def _register_default_collections(self):
         self.get(CollectionKeys.WEIGHTS).include("^(?!gradient).*weight")
         self.get(CollectionKeys.BIASES).include("^(?!gradient).*bias")
         self.get(CollectionKeys.GRADIENTS).include("^gradient")
-        self.get(CollectionKeys.LOSSES).include("[Ll]oss")
+        self.get(CollectionKeys.LOSSES).include("[Ll]oss_(?!input).*output")
 
     def create_collection(self, name):
         super().create_collection(name, cls=Collection)