Xla: Data pipeline efficiency

Created on 21 Oct 2019  路  1Comment  路  Source: pytorch/xla

Hi, and thanks for a beautiful library.

I have a few questions regarding data pipeline efficiency, originating from this discussion on Kaggle: https://www.kaggle.com/c/rsna-intracranial-hemorrhage-detection/discussion/113284

  1. Assuming a powerful enough machine on GCP, should pytorch/xla speed be similar to TF+TPU, when TPU is reading data in TFRecords format from a gs bucket (as I understand that's the most efficient setup for TF + TPU)?
  2. Are there any plans to have something like reading already prepared items from a gs bucket for pytorch/xla?
question

Most helpful comment

One part of it (reading TFRecord) is coming: https://github.com/pytorch/xla/pull/1220
TF uses infeed, which better utilize our dual VM (user and TPU) infrastructure.
The infeed ops in TF run directly on the TPU VM (and so is the input tpipeline), while in PyTorch the data is read from the user VM (processed there) and then sent to the TPU VM.
Reading directly from GCS should be possible, though in general, because of the one extra op, performance is unlikely to be as fast as TF.
Since we overlap data transmission to TPU VM, and TPU computation, the performance difference might be hidden in some case.
We are also working on simplifying the dual VM architecture, at which point there should be no more difference.

>All comments

One part of it (reading TFRecord) is coming: https://github.com/pytorch/xla/pull/1220
TF uses infeed, which better utilize our dual VM (user and TPU) infrastructure.
The infeed ops in TF run directly on the TPU VM (and so is the input tpipeline), while in PyTorch the data is read from the user VM (processed there) and then sent to the TPU VM.
Reading directly from GCS should be possible, though in general, because of the one extra op, performance is unlikely to be as fast as TF.
Since we overlap data transmission to TPU VM, and TPU computation, the performance difference might be hidden in some case.
We are also working on simplifying the dual VM architecture, at which point there should be no more difference.

Was this page helpful?
0 / 5 - 0 ratings