Skip to content
Snippets Groups Projects
Commit d046cc6e authored by Jannis Klinkenberg's avatar Jannis Klinkenberg
Browse files

fixed native tensorflow version

parent 7ecfb8bb
No related branches found
No related tags found
No related merge requests found
...@@ -66,19 +66,15 @@ def load_dataset(args): ...@@ -66,19 +66,15 @@ def load_dataset(args):
n_train, n_test = x_train.shape[0], x_test.shape[0] n_train, n_test = x_train.shape[0], x_test.shape[0]
resize_size = 224 # use bigger images with ResNet resize_size = 224 # use bigger images with ResNet
# disable any automatic data sharding in TensorFlow as we handle that manually here
# options = tf.data.Options()
# options.experimental_distribute.auto_shard_policy = tf.data.experimental.AutoShardPolicy.OFF
# Generating input pipelines # Generating input pipelines
ds_train = (tf.data.Dataset.from_tensor_slices((x_train, y_train)) ds_train = (tf.data.Dataset.from_tensor_slices((x_train, y_train))
.map(lambda image, label: (tf.image.resize(image, [resize_size, resize_size]), label)) .map(lambda image, label: (tf.image.resize(image, [resize_size, resize_size]), label))
.shuffle(n_train) # .shard(num_shards=args.world_size, index=args.world_rank) .shuffle(n_train)
.cache().batch(args.global_batch_size).prefetch(tf.data.AUTOTUNE) #.with_options(options) .cache().batch(args.global_batch_size).prefetch(tf.data.AUTOTUNE)
) )
ds_test = (tf.data.Dataset.from_tensor_slices((x_test, y_test)) ds_test = (tf.data.Dataset.from_tensor_slices((x_test, y_test))
.map(lambda image, label: (tf.image.resize(image, [resize_size, resize_size]), label)) .map(lambda image, label: (tf.image.resize(image, [resize_size, resize_size]), label))
.shuffle(n_test).cache().batch(args.global_batch_size).prefetch(tf.data.AUTOTUNE) #.with_options(options) .shuffle(n_test).cache().batch(args.global_batch_size).prefetch(tf.data.AUTOTUNE)
) )
# get updated shapes # get updated shapes
...@@ -107,7 +103,7 @@ def setup(args): ...@@ -107,7 +103,7 @@ def setup(args):
print("") print("")
sys.stdout.flush() sys.stdout.flush()
tf.config.set_visible_devices(gpu_devices[args.local_rank], "GPU") tf.config.set_visible_devices(gpu_devices[0], "GPU")
tf.keras.backend.clear_session() tf.keras.backend.clear_session()
tf.config.optimizer.set_jit(True) tf.config.optimizer.set_jit(True)
...@@ -118,6 +114,9 @@ def setup(args): ...@@ -118,6 +114,9 @@ def setup(args):
) )
) )
print("MultiWorkerMirroredStrategy.num_replicas_in_sync", strategy.num_replicas_in_sync)
print("MultiWorkerMirroredStrategy.worker_index", strategy.cluster_resolver.task_id)
return strategy return strategy
def main(): def main():
......
0% Loading or .
You are about to add 0 people to the discussion. Proceed with caution.
Please register or to comment