【pytorch】定义自己的dataloader

最新推荐文章于 2025-06-24 16:02:16 发布

唐BiuBiu

最新推荐文章于 2025-06-24 16:02:16 发布

阅读量4k

点赞数 6

CC 4.0 BY-SA版权

分类专栏：机器学习文章标签：深度学习 python

本文链接：https://blog.csdn.net/tangbiubiu/article/details/119676919

机器学习专栏收录该内容

21 篇文章

订阅专栏

本文详细介绍了如何在PyTorch中定义自定义数据加载器。首先，需要创建一个类，该类继承自`torch.utils.data.Dataset`。在类中实现`__init__`、`__getitem__`和`__len__`方法。`__init__`用于初始化数据集，包括读取数据和预处理。预处理分为训练和测试两种情况。`__getitem__`根据索引返回数据项和标签，根据训练或测试状态应用相应的预处理。`__len__`返回数据集的大小。最后，通过`torch.utils.data.DataLoader`创建数据加载器，设置批大小和是否打乱数据等参数。通过迭代加载器，可以访问到预处理后的数据。

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

在使用自己数据集训练网络时，往往需要定义自己的dataloader。这里用最简单的例子做个记录。

定义datalaoder

一般将dataloader封装为一个类，这个类继承自torch.utils.data.dataset

from torch.utils.data import dataset

class LoadData(dataset.Dataset):  # 注意父类的名称，不能写dataset
    pass

需要注意的是dataset是模块名，而Dataset是类名，在python中模块名和类名是完全独立的命名空间，因此这里的父类需要写成dataset.Dataset

在我们定义的LoadData中，至少需要有三个方法：

__init__方法，主要用来定义数据的预处理
__getitem__方法，返回数据的item和label
__len__方法，返回数据个数

from torch.utils.data import dataset

class LoadData(dataset.Dataset):
    
    def __init__(self):
        super(LoadData, self).__init__()
        pass
    
    def __getitem__(self):
        pass
    
    def __len__(self):
        pass

__init__方法需要传入至少两个参数：

一般数据的地址和标签已经被保存在某个文档中了（这里是txt格式的文档）。因此需要传入这个文档的地址。
因为__init__方法要做预处理，一般用来train的预处理和test的预处理是不同的，因此需要区分二者的参数。

	def __init__(self, txt_path, train=True):
        super(LoadData, self).__init__()
        self.img_info = self.get_img(txt_path)
        self.train = train

		# train预处理
        self.train_transforms = transforms.Compose([
            transforms.Resize(20),
            transforms.RandomHorizontalFlip(),
            transforms.RandomVerticalFlip(),
            transforms.ToTensor(),
            transforms.Normalize(mean=[0.5], std=[0.5])
        ])

		# test预处理
        self.test_transforms = transforms.Compose([
            transforms.Resize(20),
            transforms.ToTensor(),
            transforms.Normalize(mean=[0.5], std=[0.5])
        ])

	# 这个函数是用来读txt文档的
    def get_img(self, txt_path):
        with open(txt_path, 'r', encoding='utf-8') as f:
            imgs_info = f.readlines()
            imgs_info = list(map(lambda x:x.strip().split('\t'), imgs_info))
            return imgs_info

__getitem__方法只需要根据index返回数据的item和label。

    def __getitem__(self, index):
        img_path, label = self.img_info[index]
        img = Image.open(img_path)
        label = int(label)
        
        # 注意区分预处理
        if self.train:
            img = self.train_transforms(img)
        else:
            img = self.test_transforms(img)

        return img, label

__len__方法最简单，仅返回数据项个数。

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

调用dataloader

以训练数据为例，调用dataloader需要两步：

将自定义的LoadData实例化
传入torch.utils.data.dataloader中

from torch.utils.data import dataloader

train_dataset = LoadData(txt_path='XXXX', train=True)

train_loader = dataloader.Dataloader(
	dataset=train_dataset,
    batch_size=8,
    shuffle=True
	)

至此，一个最简单的dataloader就完成了！

可以用以下代码测试：

for image, label in train_loader:
	print(image.shape)
    print(label)