本文主要是介绍将自己的数据集加载到dataloader中，希望对大家解决编程问题提供一定的参考价值，需要的开发者们随着小编来一起学习吧！
from torch.utils.data import Dataset
class YourDataset(Dataset): # 继承Dataset类# 构造函数必须存在def __init__(self, root_dir, ann_file, transform=None):self.ann_file = ann_fileself.root_dir = root_dirself.img_label = self.load_annotations() # img_label是一个字典self.img = [os.path.join(self.root_dir, img) for img in list(self.img_label.keys())]self.label = [label for label in list(self.img_label.values())]self.transform = transform # 数据需要做的预处理操作def __len__(self):return len(self.img)# 获取图像和标签交给模型，该函数必须存在# 不要修改参数，每次调用时会传入随机的idx# 一个batch的数据就是由__getitem__函数处理数据传入得到的def __getitem__(self, idx):image = Image.open(self.img[idx]).convert('RGB') # img保存了图像的路径label = self.label[idx]if self.transform:image = self.transform(image) # 对数据进行预处理操作label = torch.from_numpy(np.array(label)) # 转换label的数据类型，由list->numpy->tensorreturn image, labeldef load_annotations(self):data_infos = {}with open(self.ann_file) as f:samples = [x.strip().split(' ') for x in f.readlines()]for filename, gt_label in samples:data_infos[filename] = np.array(gt_label, dtype=np.int64)return data_infos前言
使用开源算法进行小规模的训练，例如分类任务需要加载自己的数据集和类别，需要使用dataloader格式的数据集。现在通过以下的方式将自己的数据集制作为dataloader格式：
参考
参考文档：深度学习(17)--DataLoader自定义数据集制作_自定义dataloader-CSDN博客
首先感谢分享，但是文档中有较多不明白的地方，并且也没有引用，在此上做了相对的改进：
实现过程
1 从txt文件中读取图片文件名和对应label
import numpy as np
def load_annotation(ann_file): # 参数为文本文件的路径# 创建一个字典结构用于保存数据，key作为图像的名字，value作为图像的标签data_infos = {}with open(ann_file) as f:# strip()去除一些换行符等# split(' ')是以空格为分隔符# samples是一个list，格式为图像名字，图像标签# eg:[['image11.jpg,'0'],['image22.jpg,'1'],['image33.jpg,'3']]samples = [x.strip().split(' ') for x in f.readlines()]for filename, gt_label in samples:# filename是图像名字--'image11.jpg'，gt_label--'0'是标签，加载到字典data_infos中去# value值设置为array(gt_label,dtype=int64)类型data_infos[filename] = np.array(gt_label, dtype=np.int64)# 得到的字典格式：{'image11.jpg':array(0,dtype=int64),'image22.jpg':array(1,dtype=int64)}return data_infos文件格式如下：
c1.jpg 1
c2.jpg 1
c3.jpg 1
c4.jpg 1
c5.jpg 1
l1.jpg 2
l2.jpg 2
l3.jpg 2
l4.jpg 2
l5.jpg 2
l6.jpg 2
q1.jpg 3
q2.jpg 3
q3.jpg 3
q4.jpg 3
q5.jpg 3
q6.jpg 32 文件名和标签存入list
img_label = load_annotation('testload.txt')
image_name = list(img_label.keys()) # 取keys值
label = list(img_label.values()) # 取labels值
print(img_label.keys()) # 参看你的数据3 整合为数据类
from torch.utils.data import Dataset
class YourDataset(Dataset): # 继承Dataset类# 构造函数必须存在def __init__(self, root_dir, ann_file, transform=None):self.ann_file = ann_fileself.root_dir = root_dirself.img_label = self.load_annotations() # img_label是一个字典self.img = [os.path.join(self.root_dir, img) for img in list(self.img_label.keys())]self.label = [label for label in list(self.img_label.values())]self.transform = transform # 数据需要做的预处理操作def __len__(self):return len(self.img)# 获取图像和标签交给模型，该函数必须存在# 不要修改参数，每次调用时会传入随机的idx# 一个batch的数据就是由__getitem__函数处理数据传入得到的def __getitem__(self, idx):image = Image.open(self.img[idx]).convert('RGB') # img保存了图像的路径label = self.label[idx]if self.transform:image = self.transform(image) # 对数据进行预处理操作label = torch.from_numpy(np.array(label)) # 转换label的数据类型，由list->numpy->tensorreturn image, labeldef load_annotations(self):data_infos = {}with open(self.ann_file) as f:samples = [x.strip().split(' ') for x in f.readlines()]for filename, gt_label in samples:data_infos[filename] = np.array(gt_label, dtype=np.int64)return data_infos4 使用transform函数进行数据处理
# 创建一个字典结构的数据类型来进行图像预处理操作：key - value
import torchvision.transforms as transforms
data_transforms = {# 对训练集的预处理'train': transforms.Compose([transforms.Resize([256, 256]), # 卷积神经网络处理的数据大小必须相同，通过Resize来设置# 数据增强transforms.RandomRotation(45), # 随机旋转，-45到45度之间随机选#transforms.CenterCrop(64), # 从中心开始裁剪，将原本96x96大小的图片数据裁剪为64x64大小的图片数据，可以获取更多的参数transforms.RandomHorizontalFlip(p=0.5), # 随机水平翻转 选择一个概率概率，50%的概率进行水平翻转transforms.RandomVerticalFlip(p=0.5), # 随 机垂直翻转，50%的概率进行竖直翻转#transforms.ColorJitter(brightness=0.2, contrast=0.1, saturation=0.1, hue=0.1), # 参数1为亮度，参数2为对比度，参数3为饱和度，参数4为色相#transforms.RandomGrayscale(p=0.025), # 概率转换成灰度率，3通道就是R=G=B(三颜色通道转为单一颜色通道，很少进行此处理)# 将数据转为Tensor类型transforms.ToTensor(),# 标准化transforms.Normalize([0.5, 0.5, 0.5], [0.224, 0.224, 0.225]) # 设置均值，标准差，分别对应R、G、B三个颜色通道的三个均值和标准差值，(x-μ)/σ]),# 对验证集的预处理(不需要进行数据增强)'valid': transforms.Compose([transforms.Resize(256),transforms.CenterCrop(224),transforms.ToTensor(),transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])# 均值和标准差数值的设置和训练集的相同(验证集的数据对我们来说是未知的，不能利用其中的数据再计算出相关的均值和标准差)]),
}5 进行实例化
import os
from torch.utils.data import DataLoader
# 训练集
train_dataset = YourDataset(root_dir='.../...', ann_file='xxx.txt', transform=data_transforms['train'])
# 测试集
#valid_dataset = YourDataset(root_dir=valid_dir, ann_file='./.../valid.txt', transform=data_transforms['valid'])
# 实例化DataLoader(使用封装好的DataLoader包)
train_loader = DataLoader(train_dataset, batch_size=1, shuffle=True)
#valid_loader = DataLoader(valid_dataset, batch_size=64, shuffle=True)6 调用使用你的数据
dataloaders = {'train': train_loader}
#dataloaders = {'train': train_loader, "valid": valid_loader}
for inputs, labels in dataloaders['train']:print("处理训练集")
#for inputs, labels in dataloaders['valid']:
# print("处理验证集")#或者使用枚举
#for i, (imgs, labels) in enumerate(dataloaders['train']):#测试完成后你可以将以上6步写入你的程序，或封装成数据读取包7 检查和展示你的数据
# 检查训练集
from PIL import Image
import torch
import matplotlib.pyplot as plt
image1, label1 = next(iter(train_loader)) # iter表示train_loader进行迭代，next取一个batch的数据
sample = image1[0].squeeze() # 通过squeeze()压缩一个维度，有时候维度为1x3x64x64，去除这个1
# 此时的sample是3x64x64的结构，而需要图像展示则需要转换结构为64X64X3，同时需要转换为numpy数据结构
sample = sample.permute((1, 2, 0)).numpy()
# 标准化还原 x = (x-μ) / σ -> x = x*σ + μ (预处理中进行了标准化，需要还原）
sample *= [0.229, 0.224, 0.225]
sample += [0.485, 0.456, 0.406]
plt.imshow(sample)
plt.show()
print('Label is: {}'.format(label1[0].numpy()))# 检查训练集
#image2, label2 = next(iter(valid_loader)) # iter表示train_loader进行迭代，next取一个batch的数据
#sample = image2[0].squeeze() # 通过squeeze()压缩一个维度，有时候维度为1x3x64x64，去除这个1
# 此时的sample是3x64x64的结构，而需要图像展示则需要转换结构为64X64X3，同时需要转换为numpy数据结构
#sample = sample.permute((1, 2, 0)).numpy()
# 标准化还原 x = (x-μ) / σ -> x = x*σ + μ (预处理中进行了标准化，需要还原）
#sample *= [0.229, 0.224, 0.225]
#sample += [0.485, 0.456, 0.406]
#plt.imshow(sample)
#plt.show()
#print('Label is: {}'.format(label2[0].numpy()))这篇关于将自己的数据集加载到dataloader中的文章就介绍到这儿，希望我们推荐的文章对编程师们有所帮助！
