Pytorch加载数据
pytorch加载数值主要分为两个部分:
- 复写Dataset类建立一个my_dataset类用于将数据集从硬盘中逐条读取
- 利用Dataloader模块读取数据
构建Dataset:
构建dataset类主要需要几个步骤:
1.初始化def init,一般传入数据文件地址和transform
2 grtitem方法, 后续被dataloader调用的方法,主要传入index,抽出每一个样本,而且对样本的特征和标签做区分。同时还会处理传入的数据。
3 __len__用于return数据的数量,会被调用。
4自定义的一个方法 get_info,用于将数据从硬盘中读取出来,并且将数据读取出来生成一个列表,以图片为例【(图片, 对应标签)】
import os
from torch.utils.data import Dataset
from PIL import Image
label_c_d = {'cats': 0, 'dogs': 1}
class cat_dog_dataset(Dataset):
def __init__(self, data_dir, transform=None):
#data_dir中存储所有文件路径和标签
self.data_dir = self.get_img_info(data_dir)
self.label_name = {'cats': 0, 'dogs': 1}
self.transform = transform
def __getitem__(self, index):
'''
__getitem__是Dataset内部定义的方法用于读取单个
:param index: 输入索引
:return: 返回一个样本
'''
path_img, label = self.data_dir[index]
img = Image.open(path_img).convert('RGB')
if self.transform is not None:
img = self.transform(img)
return img, label
def __len__(self):
return len(self.data_dir)
@staticmethod
def get_img_info(data_dir):
'''
:param data_dir: 数据文件路径:本例中到training_set/training_set
:return: 一个包含所有数据(image_path, int(label)原组的列表
'''
data_info =list()
# os.walk:返回一个含三个值的序列,1这个序列遍历了data_dir的文件路径, 2data_dir里全部子文件的路径, 和3data_dir里全部文件的文件名
for root, dirs, _ in os.walk(data_dir):
for dir in dirs:
image_names = os.listdir(os.path.join(root, dir))
image_names = list(filter(lambda x: x.endswith('jpg'), image_names))
for i in range(len(image_names)):
image_name = image_names[i]
image_path = os.path.join(root, dir, image_name)
label = label_c_d[dir]
data_info.append((image_path, int(label)))
return data_info
Dataloader模块
Dataloader模快利用sample(采样器)对dataset返回的数据进行采样和分batch。
# 构建MyDataset实例
train_data = cat_dog_dataset(data_dir=train_dir, transform=train_transform)
valid_data = cat_dog_dataset(data_dir=valid_dir, transform=valid_transform)
# 构建DataLoder
train_loader = DataLoader(dataset=train_data, batch_size=BATCH_SIZE, shuffle=True)
valid_loader = DataLoader(dataset=valid_data, batch_size=BATCH_SIZE)