compatibility w/ TF2.1

Author: Lee Yang

examples/mnist/keras/mnist_pipeline.py

@@ -65,11 +65,9 @@ def rdd_generator():
 
   multi_worker_model.fit(x=ds, epochs=args.epochs, steps_per_epoch=max_steps_per_worker, callbacks=callbacks)
 
-  if ctx.job_name == 'chief':
-    from tensorflow_estimator.python.estimator.export import export_lib
-    export_dir = export_lib.get_timestamped_export_dir(args.export_dir)
-    tf.keras.experimental.export_saved_model(multi_worker_model, export_dir)
-    # multi_worker_model.save(args.model_dir, save_format='tf')
+  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')
 
   # terminating feed tells spark to skip processing further partitions
   tf_feed.terminate()

examples/mnist/keras/mnist_spark.py

@@ -65,11 +65,9 @@ def rdd_generator():
 
   multi_worker_model.fit(x=ds, epochs=args.epochs, steps_per_epoch=max_steps_per_worker, callbacks=callbacks)
 
-  if ctx.job_name == 'chief':
-    from tensorflow_estimator.python.estimator.export import export_lib
-    export_dir = export_lib.get_timestamped_export_dir(args.export_dir)
-    tf.keras.experimental.export_saved_model(multi_worker_model, export_dir)
-    # multi_worker_model.save(args.model_dir, save_format='tf')
+  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')
 
   # terminating feed tells spark to skip processing further partitions
   tf_feed.terminate()

examples/mnist/keras/mnist_tf.py

@@ -60,11 +60,9 @@ def build_and_compile_cnn_model():
     multi_worker_model = build_and_compile_cnn_model()
   multi_worker_model.fit(x=train_datasets, epochs=args.epochs, steps_per_epoch=args.steps_per_epoch, callbacks=callbacks)
 
-  if ctx.job_name == 'chief':
-    from tensorflow_estimator.python.estimator.export import export_lib
-    export_dir = export_lib.get_timestamped_export_dir(args.export_dir)
-    tf.keras.experimental.export_saved_model(multi_worker_model, export_dir)
-    # multi_worker_model.save(args.model_dir, save_format='tf')
+  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')
 
 
 if __name__ == '__main__':

examples/mnist/keras/mnist_tf_ds.py

@@ -86,11 +86,9 @@ def build_and_compile_cnn_model():
     multi_worker_model = build_and_compile_cnn_model()
   multi_worker_model.fit(x=train_datasets, epochs=args.epochs, steps_per_epoch=steps_per_epoch, callbacks=callbacks)
 
-  if ctx.job_name == 'chief':
-    from tensorflow_estimator.python.estimator.export import export_lib
-    export_dir = export_lib.get_timestamped_export_dir(args.export_dir)
-    tf.keras.experimental.export_saved_model(multi_worker_model, export_dir)
-    # multi_worker_model.save(args.model_dir, save_format='tf')
+  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')
 
 
 if __name__ == '__main__':

examples/segmentation/segmentation_spark.py

@@ -159,18 +159,7 @@ def unet_model(output_channels):
                             validation_steps=VALIDATION_STEPS,
                             validation_data=test_dataset)
 
-  if ctx.job_name == 'chief':
-    # Workaround for: https://github.com/tensorflow/tensorflow/issues/30251
-    print("===== saving h5py model")
-    model.save(args.model_dir + ".h5")
-    print("===== re-loading model w/o DistributionStrategy")
-    new_model = tf.keras.models.load_model(args.model_dir + ".h5")
-    print("===== exporting saved_model")
-    tf.keras.experimental.export_saved_model(new_model, args.export_dir)
-    print("===== done exporting")
-  else:
-    print("===== sleeping")
-    time.sleep(90)
+  model.save(args.export_dir, save_format='tf')
 
 
 if __name__ == '__main__':

tensorflowonspark/pipeline.py

@@ -451,6 +451,7 @@ def _transform(self, dataset):
 # global on each python worker process on the executors
 pred_fn = None           # saved_model prediction function/signature.
 pred_args = None         # args provided to the _run_model() method.  Any change will invalidate the pred_fn.
+saved_model = None
 
 
 def _run_model(iterator, args, tf_args):
@@ -471,7 +472,7 @@ def _run_model(iterator, args, tf_args):
   input_tensor_names = [tensor for col, tensor in sorted(args.input_mapping.items())]
   output_tensor_names = [tensor for tensor, col in sorted(args.output_mapping.items())]
 
-  global pred_fn, pred_args
+  global pred_fn, pred_args, saved_model
 
   # cache saved_model pred_fn to avoid reloading the model for each partition
   if not pred_fn or args != pred_args:

test/test_pipeline.py

@@ -115,12 +115,11 @@ def rdd_generator():
             return
 
       ds = tf.data.Dataset.from_generator(rdd_generator, (tf.float32, tf.float32), (tf.TensorShape([2]), tf.TensorShape([1])))
-      ds = ds.batch(args.batch_size)
-
-      # disable auto-sharding dataset
+      # disable auto-sharding since we're feeding from an RDD generator
       options = tf.data.Options()
-      options.experimental_distribute.auto_shard = False
+      options.experimental_distribute.auto_shard_policy = tf.data.experimental.AutoShardPolicy.OFF
       ds = ds.with_options(options)
+      ds = ds.batch(args.batch_size)
 
       # only train 90% of each epoch to account for uneven RDD partition sizes
       steps_per_epoch = 1000 * 0.9 // (args.batch_size * ctx.num_workers)
@@ -134,9 +133,8 @@ 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)
 
-      if ctx.job_name == 'chief' and args.export_dir:
-        print("exporting model to: {}".format(args.export_dir))
-        tf.keras.experimental.export_saved_model(model, args.export_dir)
+      print("exporting model to: {}".format(args.export_dir))
+      model.save(args.export_dir, save_format='tf')
 
       tf_feed.terminate()