Почему Pytorch медленнее Tensorflow при чтении данных

Question

Я пробовал код двух версий для итерации данных MNIST для сравнения затраченного времени.

Версия Pytorch

import os
os.environ["CUDA_VISIBLE_DEVICES"] = "0"
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

import os
os.environ['TF_CPP_MIN_LOG_LEVEL'] = '3'
os.environ["CUDA_VISIBLE_DEVICES"] = "0"
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)

И результат: TF прошло 2 секунды, а Pytorch прошло 15 секунд. Итак, почему Pytorch медленнее Tensorflow при чтении данных? Я не правильно установил? Спасибо!

Answer 1

Shijia li,

Вы все делаете правильно, это просто менее быстрый код, используемый в Pytorch (возможно, по какой-то причине).

Я просмотрел исходный код Pytorch и обнаружил следующее:

1. train_loader генерирует индексы для каждого пакета (список длиной 30000)
2. train_loader индексы для извлечения.
3. Сборщик собирает данные из набора данных, но он делает только одну запись за раз, в цикле, вызываемом строкой 44 в fetch.py ​​ module:

data = [self.dataset[idx] for idx in possibly_batched_index]

I подозреваю, что это основной источник отличий от TF - он может производить всю партию как один срез за одну оптимизированную операцию (по сравнению с Python циклом в Torch)

данные сжимаются из списка кортежей тензоров в два больших тензора collate_fn и доставляются пользователю.

Если вы хотите ускорить свой код - преобразуйте данные в тензоры (эквивалент преобразований), затем самостоятельно сгенерируйте индексы и получите фрагменты данных и целей, не вызывая загрузчик. Или упакуйте полученные тензоры в TensorDataset, который должен быть намного быстрее, чем VisionDataset (используемый для MNIST).

Кто-то из команды разработчиков Pytorch может захотеть взглянуть на это.