定义自己的dataloader

  • 1 定义datalaoder
    • 1.1 __init__
    • 1.2 __getitem__
    • 1.3 __len__
  • 2 调用dataloader
  • 参考

在使用自己数据集训练网络时，往往需要定义自己的dataloader。

1 定义datalaoder

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

from torch.utils.data import datasetclass LoadData(Dataset):# 注意父类的名称，不能写datasetpass

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

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

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

整体大致架构：

import torchfrom torch.utils.data import Datasetfrom torch.utils.data import DataLoaderclass LoadData(dDataset):def __init__(self):passdef __getitem__(self,index):passdef __len__(self):passdataset = Loaddata()train_loader = DataLoader(dataset = dataset,batch_size = 32,shuffle = Ture,num_workers=2)

1.1 init

__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

1.2 getitem

__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

1.3 len

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

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

2 调用dataloader

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

  • 将自定义的LoadData实例化
  • 传入torch.utils.data.dataloader中
from torch.utils.data import Dataloadertrain_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)

参考

https://zhuanlan.zhihu.com/p/399447239