Xla: Support for TFRecords

[henrique]

Thanks for the awesome work guys!
Is there any support, plans or possibility to use TFRecords to load data directly into the model like TF does?
Cheers

[dlibenzi]

This is more of a question for the main pytorch project, as we do not carry any TF/XLA specific data loaders.

[ailzhang]

Hey @henrique can you share some use cases that TFRecords fit into?

I'm not aware of an existing plan to add it on Pytorch or Pytorch_XLA side, but we do feature requests based on priority.

If you could make this issue a feature request with motivation/use cases, we could plan accordingly in later releases. Thanks!

[henrique]

Thanks for the quick replies!
It seems to speed up dataloading (in tf) compared to a [multi-core] pytorch dataloader. Although, I couldn't find real benchmarks to double check that, copying it the data from tfrecord to host tensor then to device seems a bit wasteful (maybe I'm wrong and that's ok?).
I also still have to check how it works with largish 3D [medical] images, as just copying them from host to device already can take some time. I hope tfrecord -> tpu will speed that up a lot, and help in distributed training.

[dlibenzi]

We now have a TFRecordReader:

https://github.com/pytorch/xla/blob/master/torch_xla/utils/tf_record_reader.py

[henrique]

Hi @dlibenzi
thanks for adding the reader!
I can loop trough a tfrecord file no problem with read_record on the cpu, from host to host

I wonder if there is a way to read from google storage directly onto the devices?
It seems like tf.dataset does that, as it even requires addition permissions in the GS bucket, and uses almost no host RAM at all (apart from speed, RAM is the main issue in regular computer vision DL, for example)
cheers

[dlibenzi]

The reader should support GCS paths (gs://) ...

[dlibenzi]

If not (but it should), you can use gcsfs:

https://github.com/pytorch/xla/blob/f43941b732ae001b2b88e7eab95dd64653c24166/torch_xla/utils/gcsfs.py#L120

[henrique]

Thanks!
My main motivation was to avoid copying stuff around.
E.g. this notebook https://www.kaggle.com/hmendonca/melanoma-neat-pytorch-lightning-native-amp takes the same time per epoch on 8 TPU's as it does on a single P100 :/

Just for the record, I don't think it reads gs (directly):
TfRecordReader('gs://rxrx1-us-central1/tfrecords/by_exp_plate_site-42/HEPG2-01_p1_s1.tfrecord').read_example()

RuntimeError: tensorflow/compiler/xla/xla_client/record_reader.cc:29 : Check failed: status == ::tensorflow::Status::OK() (Data loss: corrupted record at 0 vs. OK)
gs://rxrx1-us-central1/tfrecords/by_exp_plate_site-42/HEPG2-01_p1_s1.tfrecord offset 0

note: gs://rxrx1-us-central1/tfrecords is a public repo from a Kaggle comp

[dlibenzi]

It should work as we use the TF API to open files, which has GCS support built in.

https://github.com/pytorch/xla/blob/f43941b732ae001b2b88e7eab95dd64653c24166/third_party/xla_client/record_reader.cc#L15