Create datasets outside strategy.scope in tutorials and examples.

PiperOrigin-RevId: 259023415
Savour-yu · Jul 19, 2019 · 529a0e8 · 529a0e8
1 parent c35e1f6
commit 529a0e8
Show file tree

Hide file tree

Showing 2 changed files with 7 additions and 6 deletions.
diff --git a/tensorflow_examples/models/densenet/distributed_train.py b/tensorflow_examples/models/densenet/distributed_train.py
@@ -208,6 +208,9 @@ def main(epochs,
   devices = ['/device:GPU:{}'.format(i) for i in range(num_gpu)]
   strategy = tf.distribute.MirroredStrategy(devices)
 
+  train_dataset, test_dataset, _ = utils.create_dataset(
+      buffer_size, batch_size, data_format, data_dir)
+
   with strategy.scope():
     model = densenet.DenseNet(
         mode, growth_rate, output_classes, depth_of_model, num_of_blocks,
@@ -216,9 +219,6 @@ def main(epochs,
 
     trainer = Train(epochs, enable_function, model, batch_size, strategy)
 
-    train_dataset, test_dataset, _ = utils.create_dataset(
-        buffer_size, batch_size, data_format, data_dir)
-
     train_dist_dataset = strategy.experimental_distribute_dataset(train_dataset)
     test_dist_dataset = strategy.experimental_distribute_dataset(test_dataset)
 

diff --git a/tensorflow_examples/models/nmt_with_attention/distributed_train.py b/tensorflow_examples/models/nmt_with_attention/distributed_train.py
@@ -111,10 +111,11 @@ def main(epochs, enable_function, buffer_size, batch_size, download_path,
   strategy = tf.distribute.MirroredStrategy()
   num_replicas = strategy.num_replicas_in_sync
 
+  file_path = utils.download(download_path)
+  train_ds, test_ds, inp_lang, targ_lang = utils.create_dataset(
+      file_path, num_examples, buffer_size, batch_size)
+
   with strategy.scope():
-    file_path = utils.download(download_path)
-    train_ds, test_ds, inp_lang, targ_lang = utils.create_dataset(
-        file_path, num_examples, buffer_size, batch_size)
     vocab_inp_size = len(inp_lang.word_index) + 1
     vocab_tar_size = len(targ_lang.word_index) + 1