PyTorchの学習データの読み込みが尋常じゃなく遅い

Question

タイトル通りですが，PyTorchのライブラリを用いてデータを読み込む速度が尋常じゃなく遅いです．

# データローダを定義
transform = transforms.Compose([transforms.ToTensor(), transforms.Normalize((0.5,), (0.5,))])
trainset = torchvision.datasets.MNIST(root='./', train=True, download=True, transform=transform)
trainloader = torch.utils.data.DataLoader(trainset, batch_size=256, shuffle=True)

# Trainsetの読み込み時間の計測
t = time.time()
for x in trainset:
  pass
print(time.time() - t) # 7.56477427482605

# Dataloaderの読み込み時間の計測
t = time.time()
for x, y in trainloader:
  pass

print(time.time() - t) # 7.814661026000977

1x28x28の画像を，バッチサイズ256で60,000枚を読み込むのに8秒弱はかかり過ぎだと思うのです…．

ちなみに，環境はGoogle Colabなので，マシンの性能がネックになっていることはないと思います．
pin_memory=Trueを試したり，num_workersを適当に8として実行してみたりしましたが，殆ど変化はありませんでした．

助言をいただけると幸いです．

Comments

[quickquip]

trainsetから1つの要素を取り出すのが 133 µs かかるのが遅すぎでは? という主旨でしょうか?

[meg_]

手元の環境で実施した場合と比べて遅すぎる、という意味ですか？

[wotaku]

> quickquip様
はい，その主旨で投稿させていただきました．

> meg_様
すみません．比較対象は，
(x_train, y_train), (x_test, y_test) = datasets.mnist.load_data()
x_train = x_train.astype(np.float32) / 255
x_train = (x_train - 127.5) / 127.5
BATCH_SIZE = 32
BUFFER_SIZE = 60000
train_dataset = tf.data.Dataset.from_tensor_slices(x_train).shuffle(BUFFER_SIZE).batch(BATCH_SIZE)
という処理で生成したtrain_datasetから取り出すものです．

しかし，よく考えれば，上記の処理では最初に一旦すべてのMNISTデータを配列に格納しています．
しかし，質問文のプログラムでは，trainloaderがfor内で呼び出されるたびにデータを呼び出している?という処理のように思えます．
この理解は合っていますか? また，合っているならやはり遅くなって当然ということでしょうか?

[quickquip]

情報は質問を編集して、質問文の中に含めてください。

Answer 1

https://github.com/pytorch/vision/blob/v0.6.1/torchvision/datasets/mnist.py#L82`

純粋にtorchvision.datasets.MNISTの__getitem__の実装が、呼び出されるたびにtansformの計算をしているので、その時間かと思います。"推測するな、計測せよ"の格言からすると本当にはプロファイルを取らないといけませんが、質問のソースに続けて

preprocessed  = list(trainloader)  # ここは時間がかかる

として、計算した結果を格納してしまえば、

t = time.time()
for x, y in preprocessed:
  pass

print(time.time() - t)

と結果を読み出す方はすぐに終わります。

こちらが想定している動作なのかと想像しました。

Comments

[wotaku]

記述してくださった方法で試したところ，0.0001761…程度でループを回すことができました．
kerasで書かれたものはそもそも予めデータを格納していたために，ループ時に当然早く回せていただけなのですね．（PyTorchのローダが遅い・おかしいというのではなく，そもそも比較状況が不揃いだった）

メモリに乗り切らないデータではやはりローダを適宜ループで呼び出す他ないと思いますが，メモリに乗り切るデータ量の場合はこの方法でも良さそうです．本当にありがとうございました!!