Tensorflow数据管道：缓存到磁盘的速度很慢-如何提高评估性能？

Question

I've built a data pipeline. Pseudo code is as follows:

1. dataset ->
2. dataset = augment(dataset)
3. dataset = dataset.batch(35).prefetch(1)
4. dataset = set_from_generator(to_feed_dict(dataset)) # expensive op
5. dataset = Cache('/tmp', dataset)
6. dataset = dataset.unbatch()
7. dataset = dataset.shuffle(64).batch(256).prefetch(1)
8. to_feed_dict(dataset)

1 to 5 actions are required to generate the pretrained model outputs. I cache them as they do not change throughout epochs (pretrained model weights are not updated). 5 to 8 actions prepare the dataset for training.

Different batch sizes have to be used, as the pretrained model inputs are of a much bigger dimensionality than the outputs.

The first epoch is slow, as it has to evaluate the pretrained model on every input item to generate templates and save them to the disk. Later epochs are faster, yet they're still quite slow - I suspect the bottleneck is reading the disk cache.

What could be improved in this data pipeline to reduce the issue? Thank you!

Answer 1

1. prefetch(1) means that there will be only one element prefetched, I think you may want to have it as big as the batch size or larger.
2. After first cache you may try to put it second time but without providing a path, so it would cache some in the memory.
3. Maybe your HDD is just slow? ;)
4. Another idea is you could just manually write to compressed TFRecord after steps 1-4 and then read it with another dataset. Compressed file has lower I/O but causes higher CPU usage.