Why Pytorch is slower than Tensorflow in read data
I tried two version code for iterate MNIST data to compare the elapsed time.
Pytorch version
import os
os.environ[&quot;CUDA_VISIBLE_DEVICES&quot;] = &quot;0&quot;
import torch
from torchvision import datasets, transforms
import matplotlib.pyplot as plt
import time
train_loader = torch.utils.data.DataLoader(
    datasets.MNIST('~/data', train=True, download=True,
                   transform=transforms.Compose([
                       transforms.ToTensor(),
                       # transforms.Normalize((0.1307,), (0.3081,))
                   ])),
    batch_size=30000, shuffle=True,pin_memory=True,num_workers=4)
tic = time.time()
for epoch in range(0, 5):
    for batch_idx, (data, target) in enumerate(train_loader):
        continue
toc=time.time()
print('elapsed time:',toc-tic)

Tensorflow 2.x version
import os
os.environ['TF_CPP_MIN_LOG_LEVEL'] = '3'
os.environ[&quot;CUDA_VISIBLE_DEVICES&quot;] = &quot;0&quot;
import tensorflow as tf
from tensorflow.keras import datasets
import time
(xs, ys), (xs_, ys_) = datasets.mnist.load_data()
# print('datasets:', xs.shape, ys.shape, xs.min(), xs.max())
xs = tf.convert_to_tensor(xs, dtype=tf.float32) / 255.
db = tf.data.Dataset.from_tensor_slices((xs, ys))
db = db.batch(30000)
tic = time.time()
for epoch in range(5):
    for step, (x, y) in enumerate(db):
        continue
toc = time.time()
print('elapsed time:', toc - tic)

And the result are TF elapsed 2s and Pytorch elapsed 15s.
So, why Pytorch is slower than Tensorflow in read data?
Am I setting it wrong?
Thanks!

Shijia li,
You do everything correctly, it is just less fast code used in Pytorch (perhaps for a reason).
I looked into Pytorch source code and found the following:

train_loader generates indices for each batch(30000-long list)
train_loaderindices to fetcher.
Fetcher collects data from the dataset, but it only does it one record at a time, in cycle called by line 44 in fetch.py module:

data = [self.dataset[idx] for idx in possibly_batched_index]
I suspect that this is the main source of difference with TF - it may be producing whole batch as a single slice in one optimized operation (vs. Python cycle in Torch)

data gets compacted from list of tuples of tensors into two large tensors by collate_fn and delivered to the user.

If you want to accelerate your code - convert the data into Tensors (equivalent of transforms), then generate the indices yourself and get slices of the data and targets without calling the loader. Or pack the resulting tensors into TensorDataset which should be much faster than VisionDataset (used for MNIST).
Someone at Pytorch development team may want to take a look at this.
