定义自己的dataloader
在使用自己数据集训练网络时,往往需要定义自己的dataloader。
1 定义datalaoder
一般将dataloader封装为一个类,这个类继承自 torch.utils.data.dataset
from torch.utils.data import dataset
class LoadData(Dataset): # 注意父类的名称,不能写dataset
pass
需要注意的是dataset是模块名,而Dataset是类名,在python中模块名和类名是完全独立的命名空间,因此这里的父类需要写成 dataset.Dataset。
在我们定义的LoadData中,至少需要有三个方法:
- __init__方法,主要用来定义数据的预处理
- __getitem__方法,返回数据的item和label
- __len__方法,返回数据个数
整体大致架构:
import torch
from torch.utils.data import Dataset
from torch.utils.data import DataLoader
class LoadData(dDataset):
def __init__(self):
pass
def __getitem__(self,index):
pass
def __len__(self):
pass
dataset = 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 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)