TF Keras support for TF 1.x (aws#304)

* Introduce TF Tensor class

* rename class

* Brought in most of keras2 changes to this branch

* Add keras code

* Add keras method to collection'

* Fix naming of reduction tensors

* merge tfcollection2

* WIP

* Introduce create flag in get_hook

* addressed some comments, and fixed get_hook

* Address comments

* Address comments

* Refactored wrap_optimizer logic

* Changed default value of create_if_not_exists to False, and unset gradients on each sess begin

* Update tf in CI to 1.15

* Add check for ALL when checking if collection should be saved

* Fixing more issues with keras support

* Fixed the predict bug, as well save_all bug. Also added multiple tests

* Add reduction support

* Change log message

* Add training end test

* Cleanup and fix save of metrics

* Fix losses and metrics collection should be saved

* Add tests for regex patterns

* fix import

* Check if tensor callback already exists

* Add keras tests

* fix a reduction test, and remove duplicated test

* Fix docs and examples for distributed after wrap_optimizer change, and donot initialize writer in save_for_tensor

* Get keras working again

* Fix tests, and variable name issues

* add another test

* Add all working mirrored strategy

* All keras tests working on ec2

* Make mirrored strategy work on laptop (no GPU) as well

* Add tensorboard args to new hooks

* get session hook working for mirrored strategy

* add tests for mirrored strategy

* update test

* add test

* Fix check of isinstance

* make apt less verbose

* install tfdatasets

* Change running of tests for debugging

* Tensor device map: removed the check of strategy

* Mark more tests as slow

* trigger CI

* add pip freeze

* add pip freeze

* Upgrade pip so hopefully new keras is installed

* pin keras version

* pin keras version

* run all unit tests

* run only keras with debug

Signed-off-by: Rahul Huilgol <[email protected]>

* all unit tests

* run only keras with debug

Signed-off-by: Rahul Huilgol <[email protected]>

* run only keras with debug

Signed-off-by: Rahul Huilgol <[email protected]>

* run only keras with debug

Signed-off-by: Rahul Huilgol <[email protected]>

* empty commit

Signed-off-by: Rahul Huilgol <[email protected]>

* print step 0

Signed-off-by: Rahul Huilgol <[email protected]>

* add all tests

Signed-off-by: Rahul Huilgol <[email protected]>

* add all tests

Signed-off-by: Rahul Huilgol <[email protected]>

* rename ts_tensor

* remove kwargs

Signed-off-by: Rahul Huilgol <[email protected]>

* change where collection manager is retrieved

* Skip test

* trigger CI

* reduce logging

* dont block stdout

Author: rahul003

bin/sagemaker-containers/tensorflow/tf-train.py

@@ -4,7 +4,7 @@
 import tensorflow as tf
 import time
 import uuid
-from tornasole.tensorflow import TornasoleHook, TornasoleOptimizer, SaveConfig
+from tornasole.tensorflow import TornasoleHook, SaveConfig, get_hook
 
 parser = argparse.ArgumentParser()
 parser.add_argument("--lr", type=float, help="Learning Rate", default=0.001)
@@ -18,6 +18,13 @@
 # running in Tf estimator mode, script need to accept --model_dir parameter
 parser.add_argument("--model_dir", type=str, help="model dir", default=str(uuid.uuid4()))
 args = parser.parse_args()
+
+t = str(time.time())
+hook = TornasoleHook(
+    "s3://tornasolecodebuildtest/container_testing/ts_outputs/tf" + t,
+    save_config=SaveConfig(save_interval=10),
+)
+
 # Network definition
 with tf.name_scope("foobar"):
     x = tf.placeholder(shape=(None, 2), dtype=tf.float32)
@@ -29,15 +36,10 @@
 global_step = tf.Variable(17, name="global_step", trainable=False)
 increment_global_step_op = tf.assign(global_step, global_step + 1)
 optimizer = tf.train.AdamOptimizer(args.lr)
-optimizer = TornasoleOptimizer(optimizer)
+optimizer = get_hook().wrap_optimizer(optimizer)
 optimizer_op = optimizer.minimize(loss, global_step=increment_global_step_op)
 graph = tf.get_default_graph()
 list_of_tuples = [op.outputs for op in graph.get_operations()]
-t = str(time.time())
-hook = TornasoleHook(
-    "s3://tornasolecodebuildtest/container_testing/ts_outputs/tf" + t,
-    save_config=SaveConfig(save_interval=10),
-)
 sess = tf.train.MonitoredSession(hooks=[hook])
 for i in range(args.steps):
     x_ = np.random.random((10, 2)) * args.scale

config/buildspec.yml

@@ -12,11 +12,12 @@ phases:
     commands:
         -  . config/change_branch.sh
         - su && apt-get update
-        - apt-get install sudo
-        - sudo apt-get update
-        - sudo apt-get install unzip
+        - apt-get install sudo -qq -o=Dpkg::Use-Pty=0 # silence output: https://askubuntu.com/a/668859/724247
+        - sudo apt-get update -qq -o=Dpkg::Use-Pty=0
+        - sudo apt-get install unzip -qq -o=Dpkg::Use-Pty=0
         - cd $CODEBUILD_SRC_DIR  && chmod +x config/protoc_downloader.sh && ./config/protoc_downloader.sh
-        - pip install pytest wheel pyYaml pytest-html tensorflow==1.14.0 mxnet torch xgboost pre-commit
+        - pip install -U pip
+        - pip install -q pytest wheel pyYaml pytest-html keras==2.3.1 tensorflow==1.15.0 mxnet torch xgboost pre-commit tensorflow_datasets
         - pip uninstall -y boto3 && pip uninstall -y aiobotocore && pip uninstall -y botocore
 
   pre_build:

config/tests.sh

@@ -13,12 +13,12 @@ check_logs() {
 }
 
 run_for_framework() {
-    python -m pytest --html=$REPORT_DIR/report_$1.html --self-contained-html tests/$1
+    python -m pytest --html=$REPORT_DIR/report_$1.html -v -s --self-contained-html tests/$1
     python -m pytest --html=$REPORT_DIR/test_rules_$1.html --self-contained-html -s tests/analysis/integration_testing_rules.py::test_test_rules --mode $1 --path_to_config ./tests/analysis/config.yaml --out_dir $OUT_DIR 2>&1 | tee $REPORT_DIR/test_rules_$1.log
 }
 
 export TF_CPP_MIN_LOG_LEVEL=1
-export TORNASOLE_LOG_LEVEL=debug
+export TORNASOLE_LOG_LEVEL=info
 #export BLOCK_STDOUT=TRUE
 #export BLOCK_STDERR=FALSE
 

docs/tensorflow/README.md

@@ -77,12 +77,12 @@ If you do not specify this, it saves steps under a `GLOBAL` mode.
 ```
 hook.set_mode(ts.modes.TRAIN)
 ```
-Wrap your optimizer with TornasoleOptimizer so that
+Wrap your optimizer with wrap_optimizer so that
 Tornasole can identify your gradients and automatically
 provide these tensors as part of the `gradients` collection.
 Use this new optimizer to minimize the loss.
 ```
-optimizer = ts.TornasoleOptimizer(optimizer)
+optimizer = hook.wrap_optimizer(optimizer)
 ```
 Create a monitored session with the above hook, and use this for executing your TensorFlow job.
 ```
@@ -110,12 +110,12 @@ If you do not specify this, it saves steps under a `GLOBAL` mode.
 ```
 hook.set_mode(ts.modes.TRAIN)
 ```
-Wrap your optimizer with TornasoleOptimizer so that
+Wrap your optimizer with wrap_optimizer so that
 Tornasole can identify your gradients and automatically
 provide these tensors as part of the `gradients` collection.
 Use this new optimizer to minimize the loss.
 ```
-opt = ts.TornasoleOptimizer(opt)
+opt = hook.wrap_optimizer(opt)
 ```
 Now pass this hook to the estimator object's train, predict or evaluate methods, whichever ones you want to monitor.
 ```
@@ -410,13 +410,13 @@ hook = ts.TornasoleHook(..., include_collections = ['weights'], ...)
 
 #### Gradients
 We provide an easy way to populate the collection named `gradients` with the gradients wrt to the weights.
-This can be done by wrapping around your optimizer with `TornasoleOptimizer` as follows.
+This can be done by wrapping around your optimizer with `wrap_optimizer` as follows.
 This will also enable us to access the gradients during analysis without having to identify which tensors out of the saved ones are the gradients.
 
 ```
 import tornasole.tensorflow as ts
 ...
-opt = ts.TornasoleOptimizer(opt)
+opt = hook.wrap_optimizer(opt)
 ```
 An example for this can be seen in [this script](../../examples/tensorflow/scripts/train_imagenet_resnet_hvd.py#L738)
 Alternatively, you can refer to [customize collections](#customizing-collections) for
@@ -449,7 +449,7 @@ 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`
+above approach of wrapping your optimizer with `wrap_optimizer`
 followed by passing `optimizer_variables` in the `include_collections` parameter of the hook.
 ```
 import tornasole.tensorflow as ts
@@ -461,7 +461,7 @@ Please refer [API](api.md) for more details on using collections
 ### Customizing collections
 You can also create any other customized collection yourself.
 You can create new collections as well as modify existing collections
-(such as including gradients if you do not want to use the above `TornasoleOptimizer`)
+(such as including gradients if you do not want to use the above `wrap_optimizer`)
 #### Creating or accessing a collection
 Each collection should have a unique name (which is a string).
 You can get the collection named as `collection_name` by

docs/tensorflow/examples/distributed_training/horovod_mnist_estimator.md

@@ -13,10 +13,10 @@ import tornasole.tensorflow as ts
 ```
 **Saving gradients**
 
-We need to wrap our optimizer with TornasoleOptimizer, and use this optimizer to minimize loss.
+We need to wrap our optimizer with hook.wrap_optimizer, and use this optimizer to minimize loss.
 This will also enable us to access the gradients during analysis without having to identify which tensors out of the saved ones are the gradients.
 ```
-opt = TornasoleOptimizer(opt)
+opt = hook.wrap_optimizer(opt)
 ```
 
 

docs/tensorflow/examples/distributed_training/mirrored_strategy_mnist.md

@@ -15,10 +15,10 @@ import tornasole.tensorflow as ts
 ```
 **Saving gradients**
 
-We need to wrap our optimizer with TornasoleOptimizer, and use this optimizer to minimize loss.
+We need to wrap our optimizer with hook.wrap_optimizer, and use this optimizer to minimize loss.
 This will also enable us to access the gradients during analysis without having to identify which tensors out of the saved ones are the gradients.
 ```
-optimizer = ts.TornasoleOptimizer(optimizer)
+optimizer = hook.wrap_optimizer(optimizer)
 ```
 
 
@@ -72,7 +72,7 @@ python mirrored_strategy_mnist.py \
 --tornasole_path ~/ts_outputs/mirrored_strategy_mnist \
  --steps 5000\
  --tornasole_frequency 100\
- --reductions False
+ --reductions False\
  --save_all True
 
 ```

docs/tensorflow/examples/distributed_training/parameter_server_training/parameter_server_mnist.md

@@ -15,10 +15,10 @@ import tornasole.tensorflow as ts
 ```
 **Saving gradients**
 
-We need to wrap our optimizer with TornasoleOptimizer, and use this optimizer to minimize loss.
+We need to wrap our optimizer with hook.wrap_optimizer, and use this optimizer to minimize loss.
 This will also enable us to access the gradients during analysis without having to identify which tensors out of the saved ones are the gradients.
 ```
-optimizer = ts.TornasoleOptimizer(optimizer)
+optimizer = hook.wrap_optimizer(optimizer)
 ```
 
 

docs/tensorflow/examples/mnist.md

@@ -13,11 +13,11 @@ import tornasole.tensorflow as ts
 ```
 **Saving gradients**
 
-We need to wrap our optimizer with TornasoleOptimizer, and use this optimizer to minimize loss.
+We need to wrap our optimizer with wrap_optimizer, and use this optimizer to minimize loss.
 This will also enable us to access the gradients during analysis without having to identify which tensors out of the saved ones are the gradients.
 ```
-opt = TornasoleOptimizer(opt)
-optimizer_op = optimizer.minimize(loss, global_step=increment_global_step_op)
+opt = hook.wrap_optimizer(opt)
+optimizer_op = opt.minimize(loss, global_step=increment_global_step_op)
 ```
 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.

docs/tensorflow/examples/resnet50.md

@@ -24,13 +24,13 @@ 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.
+We need to wrap our optimizer with wrap_optimizer, and use this optimizer to minimize loss.
 This will also enable us to access the gradients during analysis without having to identify which tensors out of the saved ones are the gradients.
 ```
-opt = TornasoleOptimizer(opt)
-
 include_collections.append('gradients')
 ts.TornasoleHook(..., include_collections=include_collections, ...)
+
+opt = hook.wrap_optimizer(opt)
 ```
 Note that if include_collections is not passed to the hook,
 by default Tornasole tries to save weights, gradients and losses.

docs/tensorflow/examples/simple.md

@@ -16,13 +16,13 @@ ts.TornasoleHook(..., save_all=True, ...)
 ```
 **Saving gradients**
 
-We need to wrap our optimizer with TornasoleOptimizer, and use this optimizer to minimize loss.
+We need to wrap our optimizer with wrap_optimizer, and use this optimizer to minimize loss.
 This will also enable us to access the gradients during analysis without having to identify which tensors out of the saved ones are the gradients.
 ```
-opt = TornasoleOptimizer(opt)
-optimizer_op = optimizer.minimize(loss, global_step=increment_global_step_op)
+hook = ts.TornasoleHook(..., include_collections=[..,'gradients'], ...)
+opt = hook.wrap_optimizer(opt)
+optimizer_op = opt.minimize(loss, global_step=increment_global_step_op)
 
-ts.TornasoleHook(..., include_collections=[..,'gradients'], ...)
 ```
 **Saving losses**
 

docs/tensorflow/examples/sm_resnet50.md

@@ -20,10 +20,10 @@ include_collections.append('weights')
 ```
 **Saving gradients**
 
-We need to wrap our optimizer with TornasoleOptimizer, and use this optimizer to minimize loss.
+We need to wrap our optimizer with wrap_optimizer, and use this optimizer to minimize loss.
 This will also enable us to access the gradients during analysis without having to identify which tensors out of the saved ones are the gradients.
 ```
-opt = TornasoleOptimizer(opt)
+opt = hook.wrap_optimizer(opt)
 
 include_collections.append('gradients')
 ts.TornasoleHook(..., include_collections=include_collections, ...)

examples/tensorflow/sagemaker-notebooks/tensorflow.ipynb

@@ -82,10 +82,10 @@
     "\n",
     "```python\n",
     "import tornasole.tensorflow as ts\n",
-    "# Wrap the optimizer with Tornasole optimizer to identify gradients\n",
-    "optimizer = ts.TornasoleOptimizer(optimizer)  \n",
     "# Ask TORNASOLE to save all tensors. Note: TornasoleHook is highly configurable\n",
     "hook = ts.TornasoleHook(save_all=True)  \n",
+    "# Wrap the optimizer with Tornasole optimizer to identify gradients\n",
+    "optimizer = hook.wrap_optimizer(optimizer)  \n",
     "# pass the hook to hooks parameter of monitored session\n",
     "sess = tf.train.MonitoredSession(hooks=[hook])\n",
     "```\n",
@@ -489,9 +489,9 @@
     "```\n",
     "hook.set_mode(ts.modes.TRAIN)\n",
     "```\n",
-    "Wrap your optimizer with TornasoleOptimizer so that Tornasole can identify your gradients and automatically provide these tensors as part of the `gradients` collection. Use this new optimizer to minimize your loss during training.\n",
+    "Wrap your optimizer with wrap_optimizer so that Tornasole can identify your gradients and automatically provide these tensors as part of the `gradients` collection. Use this new optimizer to minimize your loss during training.\n",
     "```\n",
-    "optimizer = ts.TornasoleOptimizer(optimizer)\n",
+    "optimizer = hook.wrap_optimizer(optimizer)\n",
     "```\n",
     "Create a monitored session with the above hook, and use this for executing your TensorFlow job.\n",
     "```\n",
@@ -519,9 +519,9 @@
     "```\n",
     "hook.set_mode(ts.modes.TRAIN)\n",
     "```\n",
-    "Wrap your optimizer with TornasoleOptimizer so that Tornasole can identify your gradients and automatically provide these tensors as part of the `gradients` collection. Use this new optimizer to minimize your loss during training.\n",
+    "Wrap your optimizer with wrap_optimizer so that Tornasole can identify your gradients and automatically provide these tensors as part of the `gradients` collection. Use this new optimizer to minimize your loss during training.\n",
     "```\n",
-    "opt = ts.TornasoleOptimizer(opt)\n",
+    "opt = hook.wrap_optimizer(opt)\n",
     "```\n",
     "Now pass this hook to the estimator object's train, predict or evaluate methods, whichever ones you want to monitor.\n",
     "```\n",

examples/tensorflow/scripts/distributed_training/horovod_mnist_estimator.py

@@ -113,7 +113,7 @@ def cnn_model_fn(features, labels, mode):
         optimizer = hvd.DistributedOptimizer(optimizer)
 
         # Tornasole: add Tornasole Optimizer
-        optimizer = ts.TornasoleOptimizer(optimizer)
+        optimizer = ts.get_hook().wrap_optimizer(optimizer)
 
         train_op = optimizer.minimize(loss=loss, global_step=tf.train.get_global_step())
         return tf.estimator.EstimatorSpec(mode=mode, loss=loss, train_op=train_op)

examples/tensorflow/scripts/distributed_training/mirrored_strategy_mnist.py

@@ -101,7 +101,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 = ts.TornasoleOptimizer(optimizer)
+        optimizer = ts.get_hook().wrap_optimizer(optimizer)
         train_op = optimizer.minimize(loss=loss, global_step=tf.train.get_global_step())
         return tf.estimator.EstimatorSpec(mode=mode, loss=loss, train_op=train_op)
 

examples/tensorflow/scripts/distributed_training/parameter_server_training/parameter_server_mnist.py

@@ -71,7 +71,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 = ts.TornasoleOptimizer(optimizer)
+        optimizer = ts.get_hook().wrap_optimizer(optimizer)
         train_op = optimizer.minimize(loss=loss, global_step=tf.train.get_global_step())
         return tf.estimator.EstimatorSpec(mode=mode, loss=loss, train_op=train_op)
 

examples/tensorflow/scripts/keras.py

@@ -0,0 +1,80 @@
+from __future__ import absolute_import, division, print_function, unicode_literals
+
+import tensorflow_datasets as tfds
+import tensorflow as tf
+from tornasole.tensorflow import TornasoleKerasHook, get_collection
+from tornasole.core.collection import CollectionKeys
+
+tfds.disable_progress_bar()
+
+
+def train_model():
+    print(tf.__version__)
+
+    datasets, info = tfds.load(name="mnist", with_info=True, as_supervised=True)
+
+    mnist_train, mnist_test = datasets["train"], datasets["test"]
+
+    strategy = tf.distribute.MirroredStrategy()
+
+    # You can also do info.splits.total_num_examples to get the total
+    # number of examples in the dataset.
+
+    num_train_examples = info.splits["train"].num_examples
+    num_test_examples = info.splits["test"].num_examples
+
+    BUFFER_SIZE = 10000
+
+    BATCH_SIZE_PER_REPLICA = 64
+    BATCH_SIZE = BATCH_SIZE_PER_REPLICA * strategy.num_replicas_in_sync
+
+    def scale(image, label):
+        image = tf.cast(image, tf.float32)
+        image /= 255
+
+        return image, label
+
+    train_dataset = mnist_train.map(scale).cache().shuffle(BUFFER_SIZE).batch(BATCH_SIZE)
+    eval_dataset = mnist_test.map(scale).batch(BATCH_SIZE)
+
+    hook = TornasoleKerasHook(
+        out_dir="~/ts_outputs/",
+        include_collections=[
+            # CollectionKeys.WEIGHTS,
+            # CollectionKeys.GRADIENTS,
+            # CollectionKeys.OPTIMIZER_VARIABLES,
+            CollectionKeys.DEFAULT,
+            # CollectionKeys.METRICS,
+            # CollectionKeys.LOSSES,
+            # CollectionKeys.OUTPUTS,
+            # CollectionKeys.SCALARS,
+        ],
+        save_all=True,
+    )
+
+    with strategy.scope():
+        model = tf.keras.Sequential(
+            [
+                tf.keras.layers.Conv2D(32, 3, activation="relu", input_shape=(28, 28, 1)),
+                tf.keras.layers.MaxPooling2D(),
+                tf.keras.layers.Flatten(),
+                tf.keras.layers.Dense(64, activation="relu"),
+                tf.keras.layers.Dense(10, activation="softmax"),
+            ]
+        )
+        model.compile(
+            loss="sparse_categorical_crossentropy",
+            optimizer=hook.wrap_optimizer(tf.keras.optimizers.Adam()),
+            metrics=["accuracy"],
+        )
+
+    # get_collection('default').include('Relu')
+
+    callbacks = [
+        hook
+        # tf.keras.callbacks.TensorBoard(log_dir='./logs'),
+    ]
+
+    model.fit(train_dataset, epochs=1, callbacks=callbacks)
+    model.predict(eval_dataset, callbacks=callbacks)
+    model.fit(train_dataset, epochs=1, callbacks=callbacks)