pytorch中的DataLoader使用多线程读入，例子

最新推荐文章于 2025-03-29 16:53:31 发布

樱木之

最新推荐文章于 2025-03-29 16:53:31 发布

阅读量2.2k

点赞数 1

分类专栏： yuque 文章标签： pytorch 深度学习 python

本文链接：https://blog.csdn.net/qq_43369406/article/details/129967446

版权

yuque 专栏收录该内容

91 篇文章

订阅专栏

本文介绍了如何在PyTorch中利用DataLoader和Dataset类实现多线程数据加载，以提高训练效率。通过创建自定义的CustomDataset类，应用数据预处理，并设置DataLoader的num_workers参数启动多个线程，实现批量(batch_size)且随机(shuffle)地加载数据。

摘要生成于 C知道，由 DeepSeek-R1 满血版支持，前往体验 >

PyTorch中的DataLoader和Dataset可以使用多线程读取数据，这可以提高数据加载的效率。在PyTorch中，可以使用torch.utils.data.DataLoader和torch.utils.data.Dataset来实现多线程读取数据。

下面是一个简单的例子，展示如何使用多线程读取数据：

import torch
from torch.utils.data import DataLoader, Dataset
from torchvision import transforms

class CustomDataset(Dataset):
    def __init__(self, data):
        self.data = data
        self.transform = transforms.Compose([transforms.ToTensor()])

    def __getitem__(self, index):
        img = self.data[index]
        img = self.transform(img)
        return img

    def __len__(self):
        return len(self.data)

data = [img1, img2, img3, ...]
dataset = CustomDataset(data)

dataloader = DataLoader(dataset, batch_size=32, num_workers=4, shuffle=True)

在这个例子中，我们定义了一个自定义的数据集CustomDataset，其中__getitem__方法对数据进行预处理并返回预处理后的数据。然后，我们使用DataLoader将这个数据集加载进来，设置num_workers参数为4表示使用4个线程来加载数据，batch_size参数为32表示每个batch中包含32个样本，shuffle参数为True表示在每个epoch开始时打乱数据的顺序。

这样就可以使用多线程来加载数据了。注意，如果数据集很小，使用多线程加载数据可能会更慢，因为多线程有一定的开销。在这种情况下，最好使用单线程读取数据。