【pytorch】定义自己的dataloader

在使用自己数据集训练网络时，往往需要定义自己的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)