introduce compat library to support both TF2.0 and TF2.1

Author: Lee Yang

examples/mnist/estimator/README.md

@@ -1,8 +1,8 @@
 # MNIST using Estimator
 
-Original Source: https://www.tensorflow.org/beta/tutorials/distribute/multi_worker_with_estimator
+Original Source: https://www.tensorflow.org/tutorials/distribute/multi_worker_with_estimator
 
-This is the [Multi-worker Training with Estimator](https://www.tensorflow.org/beta/tutorials/distribute/multi_worker_with_estimator) example, adapted for TensorFlowOnSpark.
+This is the [Multi-worker Training with Estimator](https://www.tensorflow.org/tutorials/distribute/multi_worker_with_estimator) example, adapted for TensorFlowOnSpark.
 
 Note: this example assumes that Spark, TensorFlow, and TensorFlowOnSpark are already installed.
 

examples/mnist/estimator/mnist_spark.py

@@ -9,6 +9,8 @@ def main_fun(args, ctx):
   import tensorflow_datasets as tfds
   from tensorflowonspark import TFNode
 
+  strategy = tf.distribute.experimental.MultiWorkerMirroredStrategy()
+
   tfds.disable_progress_bar()
 
   class StopFeedHook(tf.estimator.SessionRunHook):
@@ -91,7 +93,7 @@ def model_fn(features, labels, mode):
         train_op=optimizer.minimize(
             loss, tf.compat.v1.train.get_or_create_global_step()))
 
-  strategy = tf.distribute.experimental.MultiWorkerMirroredStrategy()
+  # strategy = tf.distribute.experimental.MultiWorkerMirroredStrategy()
   config = tf.estimator.RunConfig(train_distribute=strategy, save_checkpoints_steps=100)
 
   classifier = tf.estimator.Estimator(

examples/mnist/keras/README.md

@@ -1,8 +1,8 @@
 # MNIST using Keras
 
-Original Source: https://www.tensorflow.org/beta/tutorials/distribute/multi_worker_with_keras
+Original Source: https://www.tensorflow.org/tutorials/distribute/multi_worker_with_keras
 
-This is the [Multi-worker Training with Keras](https://www.tensorflow.org/beta/tutorials/distribute/multi_worker_with_keras) example, adapted for TensorFlowOnSpark.
+This is the [Multi-worker Training with Keras](https://www.tensorflow.org/tutorials/distribute/multi_worker_with_keras) example, adapted for TensorFlowOnSpark.
 
 Notes:
 - This example assumes that Spark, TensorFlow, TensorFlow Datasets, and TensorFlowOnSpark are already installed.

examples/mnist/keras/mnist_pipeline.py

@@ -6,7 +6,7 @@
 def main_fun(args, ctx):
   import numpy as np
   import tensorflow as tf
-  from tensorflowonspark import TFNode
+  from tensorflowonspark import compat, TFNode
 
   strategy = tf.distribute.experimental.MultiWorkerMirroredStrategy()
 
@@ -67,7 +67,7 @@ def rdd_generator():
 
   from tensorflow_estimator.python.estimator.export import export_lib
   export_dir = export_lib.get_timestamped_export_dir(args.export_dir)
-  multi_worker_model.save(export_dir, save_format='tf')
+  compat.export_saved_model(multi_worker_model, export_dir, ctx.job_name == 'chief')
 
   # terminating feed tells spark to skip processing further partitions
   tf_feed.terminate()

examples/mnist/keras/mnist_spark.py

@@ -6,7 +6,7 @@
 def main_fun(args, ctx):
   import numpy as np
   import tensorflow as tf
-  from tensorflowonspark import TFNode
+  from tensorflowonspark import compat, TFNode
 
   strategy = tf.distribute.experimental.MultiWorkerMirroredStrategy()
 
@@ -67,7 +67,7 @@ def rdd_generator():
 
   from tensorflow_estimator.python.estimator.export import export_lib
   export_dir = export_lib.get_timestamped_export_dir(args.export_dir)
-  multi_worker_model.save(export_dir, save_format='tf')
+  compat.export_saved_model(multi_worker_model, export_dir, ctx.job_name == 'chief')
 
   # terminating feed tells spark to skip processing further partitions
   tf_feed.terminate()

examples/mnist/keras/mnist_tf.py

@@ -6,6 +6,8 @@
 def main_fun(args, ctx):
   import tensorflow_datasets as tfds
   import tensorflow as tf
+  from tensorflowonspark import compat
+
   tfds.disable_progress_bar()
 
   strategy = tf.distribute.experimental.MultiWorkerMirroredStrategy()
@@ -62,7 +64,7 @@ def build_and_compile_cnn_model():
 
   from tensorflow_estimator.python.estimator.export import export_lib
   export_dir = export_lib.get_timestamped_export_dir(args.export_dir)
-  multi_worker_model.save(export_dir, save_format='tf')
+  compat.export_saved_model(multi_worker_model, export_dir, ctx.job_name == 'chief')
 
 
 if __name__ == '__main__':

examples/mnist/keras/mnist_tf_ds.py

@@ -6,6 +6,7 @@
 def main_fun(args, ctx):
   """Example demonstrating loading TFRecords directly from disk (e.g. HDFS) without tensorflow_datasets."""
   import tensorflow as tf
+  from tensorflowonspark import compat
 
   strategy = tf.distribute.experimental.MultiWorkerMirroredStrategy()
 
@@ -88,7 +89,7 @@ def build_and_compile_cnn_model():
 
   from tensorflow_estimator.python.estimator.export import export_lib
   export_dir = export_lib.get_timestamped_export_dir(args.export_dir)
-  multi_worker_model.save(export_dir, save_format='tf')
+  compat.export_saved_model(multi_worker_model, export_dir, ctx.job_name == 'chief')
 
 
 if __name__ == '__main__':

examples/segmentation/README.md

@@ -1,8 +1,8 @@
 # Image Segmentation
 
-Original Source: https://www.tensorflow.org/beta/tutorials/images/segmentation
+Original Source: https://www.tensorflow.org/tutorials/images/segmentation
 
-This code is based on the [Image Segmentation](https://www.tensorflow.org/beta/tutorials/images/segmentation) notebook example, converted to a single-node TensorFlow python app, then converted into a distributed TensorFlow app using the `MultiWorkerMirroredStrategy`, and then finally adapted for TensorFlowOnSpark.  Compare the different versions to see the conversion steps involved at each stage.
+This code is based on the [Image Segmentation](https://www.tensorflow.org/tutorials/images/segmentation) notebook example, converted to a single-node TensorFlow python app, then converted into a distributed TensorFlow app using the `MultiWorkerMirroredStrategy`, and then finally adapted for TensorFlowOnSpark.  Compare the different versions to see the conversion steps involved at each stage.
 
 Notes: 
 - this example assumes that Spark, TensorFlow, and TensorFlowOnSpark are already installed.

examples/segmentation/segmentation_spark.py

@@ -159,7 +159,15 @@ def unet_model(output_channels):
                             validation_steps=VALIDATION_STEPS,
                             validation_data=test_dataset)
 
-  model.save(args.export_dir, save_format='tf')
+  if tf.__version__ == '2.0.0':
+    # Workaround for: https://github.com/tensorflow/tensorflow/issues/30251
+    # Save model locally as h5py and reload it w/o distribution strategy
+    if ctx.job_name == 'chief':
+      model.save(args.model_dir + ".h5")
+      new_model = tf.keras.models.load_model(args.model_dir + ".h5")
+      tf.keras.experimental.export_saved_model(new_model, args.export_dir)
+  else:
+    model.save(args.export_dir, save_format='tf')
 
 
 if __name__ == '__main__':

test/test_pipeline.py

@@ -4,6 +4,7 @@
 import test
 import unittest
 
+from tensorflowonspark import compat
 from tensorflowonspark.pipeline import HasBatchSize, HasSteps, Namespace, TFEstimator, TFParams
 from tensorflow.keras import Sequential
 from tensorflow.keras.layers import Dense
@@ -117,7 +118,7 @@ def rdd_generator():
       ds = tf.data.Dataset.from_generator(rdd_generator, (tf.float32, tf.float32), (tf.TensorShape([2]), tf.TensorShape([1])))
       # disable auto-sharding since we're feeding from an RDD generator
       options = tf.data.Options()
-      options.experimental_distribute.auto_shard_policy = tf.data.experimental.AutoShardPolicy.OFF
+      compat.disable_auto_shard(options)
       ds = ds.with_options(options)
       ds = ds.batch(args.batch_size)
 
@@ -133,8 +134,9 @@ def rdd_generator():
       # This fails with: "NotImplementedError: `fit_generator` is not supported for models compiled with tf.distribute.Strategy"
       # model.fit_generator(ds, epochs=args.epochs, steps_per_epoch=steps_per_epoch, callbacks=callbacks)
 
-      print("exporting model to: {}".format(args.export_dir))
-      model.save(args.export_dir, save_format='tf')
+      if args.export_dir:
+        print("exporting model to: {}".format(args.export_dir))
+        compat.export_saved_model(model, args.export_dir, ctx.job_name == 'chief')
 
       tf_feed.terminate()