From 37a04fd6df4bf89f5bcb9f6a0bce7c02ff127181 Mon Sep 17 00:00:00 2001 From: Hunor Laczko Date: Wed, 25 May 2022 23:02:56 +0200 Subject: [PATCH] add dataset options --- train.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/train.py b/train.py index 20d2a96..83b82d3 100644 --- a/train.py +++ b/train.py @@ -401,7 +401,7 @@ def main(): # prepare dataset # dataset = get_ffhq_dataset(args['tfrecord_dir'], args['train_res'], batch_size=global_batch_size, epochs=None) - dataset = DatasetSimple('/home/worksapce/data/dataset/stylegan').batch(global_batch_size)#.shuffle(buffer_size=1000).prefetch(buffer_size=tf.data.experimental.AUTOTUNE).cache() + dataset = DatasetSimple('/workspace/Dataset/stylegan').batch(global_batch_size).shuffle(buffer_size=1000).prefetch(buffer_size=tf.data.experimental.AUTOTUNE).cache().repeat() options = tf.data.Options() options.experimental_distribute.auto_shard_policy = tf.data.experimental.AutoShardPolicy.OFF dataset = dataset.with_options(options)