Regarding performance of petastorm with shuffle

Hi, I have 24GB data set (1.28 million images), stored as parquet files with 800 mb each with total as 31 parquet files. I have 60 core cpu machine with 120GB RAM.

I use below code

def just_read_batch(dataset_url):
    with make_batch_reader(dataset_url, num_epochs=1, workers_count=16, shuffle_row_drop_partitions=10) as train_reader:
        i = 0
        for schema_view in train_reader:
            i += len(schema_view.imagename)
        print(i)

parquet with shuffle (shuffle_row_drop_partitions = 5) : 202 seconds (PERFORMANCE IS 6k image per second) run again with shuffle : 200 seconds run again with shuffle : 200 seconds

with (shuffle_row_drop_partitions = 10) : 361 seconds run again : 370 seconds

how can I improve the performance ? I need shuffled data on each epoch. please check my code snippet again and suggest me whether I need to change configuration to get better performance ?

[selitvin]

Shorter version

• Don’t set shuffle_row_drop_partitions: I would expect you to get 30K images per second, unless network/disk saturates.
• If your images are encoded using Petastorm codec, using make_batch_reader would not decode them for you at the moment (a bug : you would get unusable binary blobs). You should use make_reader which will handle the decoding correctly. Decoding the images however are likely to stress your CPUs and you may see lower throughput.
• Place a shuffling queue after the reader to shuffle your data.

Longer version

shuffle_row_drop_partitions is a very wasteful parameter: it decreases correlation between produced samples at the expense of network traffic. In case of shuffle_row_drop_partitions=10, 1/10 of the downloaded samples in a rowgroup are used. (the same rowgroup would be loaded later and another 1/10 will be used next time, and so on). I’d suggest to use this switch only in extreme cases due to its wastefulness.

Reading from Apache Parquet format can be done only in atomic units of a row-group. shuffle_row_groups=True results in random order of row-groups being consumed, but there is still inter-row-group correlation issue. Here are some ideas of reducing correlation between samples:

• Placing a shuffling queue after the reader. Usually we rely on Tensorflow shuffling queue (tf_tensors take shuffling_queue_capacity parameter, or you can use tf.data facilities). I think you can also use RandomShufflingBuffer in petastorm.reader_impl.shuffling_buffer - we did not intend it to be a public class, but I think it would work as a pure-Python shuffling queue implementation.
• Reduce size of rowgroup (hence decrease amount of correlated samples). If you generate your dataset using Petastorm, materialize_dataset takes row_group_size_mb. Reducing its value will decrease number of rows in a rowgroup (note that with an off-the-shelf parquet Java library implementation you can not get less then 100 rows. We have a patch I can share with you that would allow to overcome that issue)
• Depends on the nature of data you are using, shuffling the rows in the dataset before writing it out may help (we preshuffle this with some of our datasets).

I am a bit confused with the fact that you use make_batch_reader: you mentioned you work with images, but for images, I would expect you to use Tensor and have your own custom Unischema. make_batch_reader however is designed to work with standard Apache Parquet types (I think the image fields are currently unusable in this setup (we will fix it soon)).

[selitvin]

If I understand correctly, the 2-pass shuffling mentioned in the article describes an algorithm for shuffling that utilizes additional disk storage. A preprocessing step that reshuffles and persists shuffled data is definitely something that can be done be a user but is not in the scope of the responsibility of the Petastorm library…