TFRecords for Pytorch

If you are using the dataset often, I would suggest extracting it once and saving it in another format, this way you can have fast data reads when training.

I personnaly used LMDB for a larger-than-memory sized dataset, and data loading is fast enough to saturate the GPU(s) at all times (as advised by @smth in this post BTW).