Add losses support to Tensorflow (aws#135)

* add loss support

* remove accidentally included file

* add loss not decreasing rule

* add loss rule and test

* Added integration tests

* fix tests

* fix doc

* fix test

* fix test

* fix test

* fix test

* stdout level

* fix train

* fix integration test

* bring back correct logging behavior

* add log

* fix log name

* rename job log

* address review

* fix integration test

* change behavior of required tensors, and simplify loss rule

* introduce a class for required tensors and use old one as for each trial

* undo changes of analysis

* add integraion tests

* remove new file

* remove loss rule

* remove test

* change utils method

* add losses to sagemaker doc as well

* fix loss test

Author: rahul003

docs/tensorflow/README.md

@@ -431,6 +431,25 @@ Then, you need to pass `gradients` in the `include_collections` parameter of the
 import tornasole.tensorflow as ts
 hook = ts.TornasoleHook(..., include_collections = ['gradients'], ...)
 ```
+#### Losses
+If you are using the default loss functions in Tensorflow, Tornasole can automatically pick up these losses from Tensorflow's losses collection.
+In such a case, we only need to specify 'losses' in the `include_collections` argument of the hook.
+If you do not pass this argument to the hook, it will save losses by default.
+If you are using your custom loss function, you can either add this to Tensorflow's losses collection or Tornasole's losses collection as follows:
+
+```
+import tornasole.tensorflow as ts
+
+# if your loss function is not a default TF loss function, 
+# but is a custom loss function
+# then add to the collection losses
+loss = ...
+ts.add_to_collection('losses', loss)
+
+# specify losses in include_collections 
+# Note that this is included by default
+hook = ts.TornasoleHook(..., include_collections = ['losses'..], ...)
+```
 
 #### Optimizer Variables
 Optimizer variables such as momentum can also be saved easily with the 

docs/tensorflow/examples/mnist.md

@@ -18,9 +18,17 @@ This will also enable us to access the gradients during analysis without having
 ```
 opt = TornasoleOptimizer(opt)
 optimizer_op = optimizer.minimize(loss, global_step=increment_global_step_op)
-
-ts.TornasoleHook(..., include_collections=[..,'gradients'], ...)
 ```
+Note that here since by default Tornasole tries to save weights, gradients and losses 
+we didn't need to specify 'gradients' in the include_collections argument of the hook.
+
+**Saving losses**
+
+Since we use a default loss function from Tensorflow here, 
+we would only need to indicate to the hook that we want to include losses. 
+But since the hook by default saves losses if include_collections argument was not set, 
+we need not do anything.
+
 **Setting save interval**
 
 You can set different save intervals for different modes. 

docs/tensorflow/examples/resnet50.md

@@ -17,7 +17,11 @@ import tornasole.tensorflow as ts
 **Saving weights**
 ```
 include_collections.append('weights')
+ts.TornasoleHook(..., include_collections=include_collections, ...)
 ```
+Note that the above line of include_collections is not required 
+because by default Tornasole tries to save weights, gradients and losses.
+
 **Saving gradients**
 
 We need to wrap our optimizer with TornasoleOptimizer, and use this optimizer to minimize loss. 
@@ -28,6 +32,18 @@ opt = TornasoleOptimizer(opt)
 include_collections.append('gradients')
 ts.TornasoleHook(..., include_collections=include_collections, ...)
 ```
+Note that if include_collections is not passed to the hook, 
+by default Tornasole tries to save weights, gradients and losses.
+
+**Saving losses**
+
+Since we use a default loss function from Tensorflow, we only need to indicate to the hook that we want to include losses.
+In the code, you will see the following line to do so. 
+```
+include_collections=['losses']
+ts.TornasoleHook(..., include_collections=include_collections, ...)
+```
+
 **Saving relu activations by variable**
 ```
 x = tf.nn.relu(x + shortcut)

docs/tensorflow/examples/simple.md

@@ -24,6 +24,18 @@ optimizer_op = optimizer.minimize(loss, global_step=increment_global_step_op)
 
 ts.TornasoleHook(..., include_collections=[..,'gradients'], ...)
 ```
+**Saving losses**
+
+Since we are not using a default loss function from Tensorflow, 
+we need to tell Tornasole to add our loss to the losses collection as follows
+```
+ts.add_to_collection('losses', loss)
+``` 
+In the code, you will see the following line to do so. 
+```
+ts.TornasoleHook(..., include_collections=[...,'losses'], ...)
+```
+
 **Setting save interval**
 ```
 ts.TornasoleHook(...,save_config=ts.SaveConfig(save_interval=args.tornasole_frequency)...)

examples/tensorflow/scripts/mnist.py

@@ -1,5 +1,6 @@
 import argparse
 import numpy as np
+import random
 import tensorflow as tf
 import tornasole.tensorflow as ts
 
@@ -9,6 +10,8 @@
                     help="How often to save TS data", default=50)
 parser.add_argument('--tornasole_eval_frequency', type=int,
                     help="How often to save TS data", default=10)
+parser.add_argument('--lr', type=float, default=0.001)
+parser.add_argument('--random_seed', type=bool, default=False)
 parser.add_argument('--num_epochs', type=int, default=5,
                     help="Number of epochs to train for")
 parser.add_argument('--num_steps', type=int,
@@ -17,6 +20,14 @@
 parser.add_argument('--model_dir', type=str, default='/tmp/mnist_model')
 args = parser.parse_args()
 
+# these random seeds are only intended for test purpose.
+# for now, 2,2,12 could promise no assert failure when running tornasole_rules test_rules.py with config.yaml
+# if you wish to change the number, notice that certain steps' tensor value may be capable of variation
+if args.random_seed:
+    tf.set_random_seed(2)
+    np.random.seed(2)
+    random.seed(12)
+
 def cnn_model_fn(features, labels, mode):
   """Model function for CNN."""
   # Input Layer
@@ -67,7 +78,7 @@ def cnn_model_fn(features, labels, mode):
 
   # Configure the Training Op (for TRAIN mode)
   if mode == tf.estimator.ModeKeys.TRAIN:
-    optimizer = tf.train.GradientDescentOptimizer(learning_rate=0.001)
+    optimizer = tf.train.GradientDescentOptimizer(learning_rate=args.lr)
     optimizer = ts.TornasoleOptimizer(optimizer)
     train_op = optimizer.minimize(
       loss=loss,

examples/tensorflow/scripts/simple.py

@@ -37,6 +37,7 @@
     w0 = [[1], [1.]]
     y = tf.matmul(x, w0)
 loss = tf.reduce_mean((tf.matmul(x, w) - y) ** 2, name="loss")
+ts.add_to_collection('losses', loss)
 
 global_step = tf.Variable(17, name="global_step", trainable=False)
 increment_global_step_op = tf.assign(global_step, global_step+1)
@@ -56,7 +57,7 @@
 # Note that we are saving all tensors here by passing save_all=True
 hook = ts.TornasoleHook(out_dir=args.tornasole_path,
                         save_all=True,
-                        include_collections=['weights', 'gradients'],
+                        include_collections=['weights', 'gradients', 'losses'],
                         save_config=ts.SaveConfig(save_interval=args.tornasole_frequency),
                         reduction_config=rdnc)
 

examples/tensorflow/scripts/train_imagenet_resnet_hvd.py

@@ -1020,7 +1020,7 @@ def get_tornasole_hook(FLAGS):
     else:
         rnc = None
 
-    include_collections = []
+    include_collections = ['losses']
 
     if FLAGS.tornasole_save_weights is True:
         include_collections.append('weights')

sagemaker-docs/DeveloperGuide_TF.md

@@ -304,6 +304,25 @@ import tornasole.tensorflow as ts
 hook = ts.TornasoleHook(..., include_collections = ['gradients'], ...)
 ```
 
+#### Losses
+If you are using the default loss functions in Tensorflow, Tornasole can automatically pick up these losses from Tensorflow's losses collection.
+In such a case, we only need to specify 'losses' in the `include_collections` argument of the hook.
+If you do not pass this argument to the hook, it will save losses by default.
+If you are using your custom loss function, you can either add this to Tensorflow's losses collection or Tornasole's losses collection as follows:
+```
+import tornasole.tensorflow as ts
+
+# if your loss function is not a default TF loss function, 
+# but is a custom loss function
+# then add to the collection losses
+loss = ...
+ts.add_to_collection('losses', loss)
+
+# specify losses in include_collections 
+# Note that this is included by default
+hook = ts.TornasoleHook(..., include_collections = ['losses'..], ...)
+```
+
 #### Optimizer Variables
 Optimizer variables such as momentum can also be saved easily with the 
 above approach of wrapping your optimizer with `TornasoleOptimizer` 

tests/analysis/config.yaml

@@ -6,6 +6,8 @@
   - s3: True # run test cases in s3 mode
   - [simple.py: &simple
       $CODEBUILD_SRC_DIR/examples/tensorflow/scripts/simple.py,
+     tf_mnist.py: &tf_mnist
+      $CODEBUILD_SRC_DIR/examples/tensorflow/scripts/mnist.py,
      torch_simple.py: &torch_simple
       $CODEBUILD_SRC_DIR/examples/pytorch/scripts/simple.py,
      train_imagenet_resnet_hvd.py: &train_imagenet_resnet_hvd
@@ -82,7 +84,24 @@
      *invoker,
      --rule_name weightupdateratio --flag True --end_step 71
     ]
-
+-
+  - loss_not_decreasing/tf/true
+  - tensorflow
+  - *Enable
+  - [*tf_mnist,
+     --lr 0.001 --tornasole_train_frequency 10 --random_seed True,
+     *invoker,
+     --rule_name lossnotdecreasing --flag True --end_step 1000
+    ]
+-
+  - loss_not_decreasing/tf/false
+  - tensorflow
+  - *Enable
+  - [*simple,
+     --lr 0.05 --scale 1 --steps 1009 --tornasole_frequency 13 --random_seed True,
+     *invoker,
+     --rule_name lossnotdecreasing --flag False --num_steps 100 --min_difference 12
+    ]
 
 # test cases for mxnet
 -

tests/analysis/utils.py

@@ -18,9 +18,9 @@ def generate_data(path, trial, step, tname_prefix,
     if export_colls:
         c = CollectionManager()
         c.add("default")
-        c.get("default").tensor_names = [ tname_prefix + '_' + str(i) for i in range(num_tensors)]
+        c.get("default").tensor_names = [f'{tname_prefix}_{i}' for i in range(num_tensors)]
         c.add('gradients')
-        c.get("gradients").tensor_names = [ tname_prefix + '_' + str(i) for i in range(num_tensors)]
+        c.get("gradients").tensor_names = [f'{tname_prefix}_{i}' for i in range(num_tensors)]
         c.export(os.path.join(path, trial, "collections.ts"))
 
 

tests/tensorflow/hooks/test_estimator_modes.py

@@ -138,10 +138,9 @@ def test_mnist_local():
     assert len(tr.available_steps()) == 55
     assert len(tr.available_steps(mode=ts.modes.TRAIN)) == 15
     assert len(tr.available_steps(mode=ts.modes.EVAL)) == 40
-    assert len(tr.tensors()) == 16
+    assert len(tr.tensors()) == 17
     shutil.rmtree(trial_dir)
 
-
 def test_mnist_local_json():
     out_dir = 'newlogsRunTest1/test_mnist_local_json_config'
     shutil.rmtree(out_dir, ignore_errors=True)
@@ -151,52 +150,47 @@ def test_mnist_local_json():
     assert len(tr.available_steps()) == 55
     assert len(tr.available_steps(mode=ts.modes.TRAIN)) == 15
     assert len(tr.available_steps(mode=ts.modes.EVAL)) == 40
-    assert len(tr.tensors()) == 16
+    assert len(tr.tensors()) == 17
     shutil.rmtree(out_dir, ignore_errors=True)
 
-
 def test_mnist_s3():
     run_id = 'trial_' + datetime.now().strftime('%Y%m%d-%H%M%S%f')
     trial_dir = 's3://tornasole-testing/tornasole_tf/hooks/estimator_modes/' + run_id
     tr = help_test_mnist(trial_dir, ts.SaveConfig(save_interval=2))
     assert len(tr.available_steps()) == 55
     assert len(tr.available_steps(mode=ts.modes.TRAIN)) == 15
     assert len(tr.available_steps(mode=ts.modes.EVAL)) == 40
-    assert len(tr.tensors()) == 16
-
+    assert len(tr.tensors()) == 17
 
 def test_mnist_local_multi_save_configs():
     run_id = 'trial_' + datetime.now().strftime('%Y%m%d-%H%M%S%f')
     trial_dir = os.path.join(TORNASOLE_TF_HOOK_TESTS_DIR, run_id)
     tr = help_test_mnist(trial_dir, {ts.modes.TRAIN: ts.SaveConfig(save_interval=2),
-                                     ts.modes.EVAL: ts.SaveConfig(save_interval=1)})
+                              ts.modes.EVAL: ts.SaveConfig(save_interval=1)})
     assert len(tr.available_steps()) == 94
     assert len(tr.available_steps(mode=ts.modes.TRAIN)) == 15
     assert len(tr.available_steps(mode=ts.modes.EVAL)) == 79
-    assert len(tr.tensors()) == 16
+    assert len(tr.tensors()) == 17
     shutil.rmtree(trial_dir)
 
-
 def test_mnist_s3_multi_save_configs():
     run_id = 'trial_' + datetime.now().strftime('%Y%m%d-%H%M%S%f')
     trial_dir = 's3://tornasole-testing/tornasole_tf/hooks/estimator_modes/' + run_id
     tr = help_test_mnist(trial_dir, {ts.modes.TRAIN: ts.SaveConfig(save_interval=2),
-                                     ts.modes.EVAL: ts.SaveConfig(save_interval=1)})
+                              ts.modes.EVAL: ts.SaveConfig(save_interval=1)})
     assert len(tr.available_steps()) == 94
     assert len(tr.available_steps(mode=ts.modes.TRAIN)) == 15
     assert len(tr.available_steps(mode=ts.modes.EVAL)) == 79
-    assert len(tr.tensors()) == 16
-
+    assert len(tr.tensors()) == 17
 
 def test_mnist_local_multi_save_configs_json():
     out_dir = 'newlogsRunTest1/test_save_config_modes_hook_config'
     shutil.rmtree(out_dir, ignore_errors=True)
-    os.environ[
-        TORNASOLE_CONFIG_FILE_PATH_ENV_STR] = 'tests/tensorflow/hooks/test_json_configs/test_save_config_modes_hook_config.json'
+    os.environ[TORNASOLE_CONFIG_FILE_PATH_ENV_STR] = 'tests/tensorflow/hooks/test_json_configs/test_save_config_modes_hook_config.json'
     hook = ts.TornasoleHook.hook_from_config()
     tr = help_test_mnist(out_dir, hook=hook)
     assert len(tr.available_steps()) == 94
     assert len(tr.available_steps(mode=ts.modes.TRAIN)) == 15
     assert len(tr.available_steps(mode=ts.modes.EVAL)) == 79
-    assert len(tr.tensors()) == 16
+    assert len(tr.tensors()) == 17
     shutil.rmtree(out_dir)

tests/tensorflow/hooks/test_losses.py

@@ -0,0 +1,14 @@
+from .utils import *
+import tornasole.tensorflow as ts
+import shutil
+
+from .test_estimator_modes import help_test_mnist
+
+def test_mnist_local():
+  run_id = 'trial_' + datetime.now().strftime('%Y%m%d-%H%M%S%f')
+  trial_dir = os.path.join(TORNASOLE_TF_HOOK_TESTS_DIR, run_id)
+  tr = help_test_mnist(trial_dir, ts.SaveConfig(save_interval=2))
+  assert len(tr.collection('losses').get_tensor_names()) == 1
+  for t in tr.collection('losses').get_tensor_names():
+    assert len(tr.tensor(t).steps()) == 55
+  shutil.rmtree(trial_dir)

tests/tensorflow/hooks/test_reductions.py

@@ -15,7 +15,7 @@ def helper_test_reductions(trial_dir, hook):
     from tornasole.trials import create_trial
 
     tr = create_trial(trial_dir)
-    assert len(tr.tensors()) == 2
+    assert len(tr.tensors()) == 3
     for tname in tr.tensors():
         t = tr.tensor(tname)
         try:

tests/tensorflow/hooks/test_save_all_full.py

@@ -1,5 +1,5 @@
 from .utils import *
-from tornasole.tensorflow import reset_collections, get_collections
+from tornasole.tensorflow import reset_collections, get_collections, CollectionManager, Collection
 import shutil, glob
 from tornasole.core.reader import FileReader
 from tornasole.core.json_config import TORNASOLE_CONFIG_FILE_PATH_ENV_STR
@@ -23,16 +23,18 @@ def test_save_all_full(hook=None, trial_dir=None):
     dirs, _ = get_dirs_files(os.path.join(trial_dir, 'events'))
 
     coll = get_collections()
-    assert len(coll) == 5
+    assert len(coll) == 6
     assert len(coll['weights'].tensor_names) == 1
     assert len(coll['gradients'].tensor_names) == 1
+    assert len(coll['losses'].tensor_names) == 1
 
     assert 'collections.ts' in files
     cm = CollectionManager.load(join(trial_dir, 'collections.ts'))
 
-    assert len(cm.collections) == 5
+    assert len(cm.collections) == 6
     assert len(cm.collections['weights'].tensor_names) == 1
     assert len(cm.collections['weights'].reduction_tensor_names) == 0
+    assert len(cm.collections['losses'].tensor_names) == 1
     assert len(cm.collections['gradients'].tensor_names) == 1
     assert len(cm.collections['gradients'].reduction_tensor_names) == 0
     # as we hadn't asked to be saved

tests/tensorflow/hooks/test_save_reductions.py

@@ -1,5 +1,5 @@
 from .utils import *
-from tornasole.tensorflow import reset_collections, get_collections
+from tornasole.tensorflow import reset_collections, get_collections, CollectionManager
 import shutil
 import glob
 from tornasole.core.reader import FileReader
@@ -10,13 +10,13 @@ def helper_save_reductions(trial_dir, hook):
   _, files = get_dirs_files(trial_dir)
   coll = get_collections()
 
-  assert len(coll) == 4
+  assert len(coll) == 5
   assert len(coll['weights'].reduction_tensor_names) == 1
   assert len(coll['gradients'].reduction_tensor_names) == 1
 
   assert 'collections.ts' in files
   cm = CollectionManager.load(join(trial_dir, 'collections.ts'))
-  assert len(cm.collections) == 4
+  assert len(cm.collections) == 5
   assert len(cm.collections['weights'].tensor_names) == 0
   assert len(cm.collections['weights'].reduction_tensor_names) == 1
   assert len(cm.collections['gradients'].tensor_names) == 0
@@ -45,8 +45,8 @@ def helper_save_reductions(trial_dir, hook):
         tensor_name, step, tensor_data, mode, mode_step = x
         i += 1
         size += tensor_data.nbytes if tensor_data is not None else 0
-    assert i == 32
-    assert size == 128
+    assert i == 48
+    assert size == 192
 
   shutil.rmtree(trial_dir)