[pytorch] 定义自己的dataloader

本文主要是介绍[pytorch] 定义自己的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 torch
from torch.utils.data import Dataset
from 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

这篇关于[pytorch] 定义自己的dataloader的文章就介绍到这儿，希望我们推荐的文章对编程师们有所帮助！

---