add dataset options

This commit is contained in:
2022-05-25 23:02:56 +02:00
parent 564b9e96a7
commit 37a04fd6df
+1 -1
View File
@@ -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)