CPU上keras和pytorch训练速度比较(问题已解决）

参考github issue
https://github.com/pytorch/pytorch/issues/42402

测试发现，keras在比较简单的网络上远比pytorch跑的要快的多，目前大致的猜测原因在读取数据部分，但是我还不知道怎么pytorch这个数据的读取速度，为什么都没有大佬发现这个问题呢？

mnist数据集读取

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('./', train=True, download=False,
                   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)

结果如下
而对于keras，

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)

结果如下

keras是1s, 而pytorch用了大概17秒，这也太离谱了吧，不知道是哪里出了问题，

经过大佬的修改代码，pytorch代码应该修改如下，结果读取数据只用了2s

from torch.utils.data import DataLoader
from torchvision import datasets
import time


class NoPIL_MNIST(datasets.MNIST):
    def __getitem__(self, index):
        img, target = self.data[index], int(self.targets[index])

        if self.transform is not None:
            img = self.transform(img)

        if self.target_transform is not None:
            target = self.target_transform(target)

        return img, target


if __name__ == "__main__":
    dataset = NoPIL_MNIST(
        './',
        train=True, download=True,
    )

    train_loader = DataLoader(
        dataset,
        batch_size=30000,
        shuffle=True,
        pin_memory=True,
        num_workers=0
    )

    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)