python读取文件数据生成图像_pytorch ImageFolder和Dataloader加载自制图像数据集

0 总结

首先对数据集进行文件夹组织处理以适应ImageFolder的要求，然后将ImageFolder的返回结果作为DataLoader的输入，完成训练集数据和测试集数据的迭代。

1.数据集简介

数据集是从网上下载的宝可梦图像数据集，下载后的数据集已按文件夹名称分为5类，每类大约有220~240张图像，格式有jpg, jpeg, png和gif，经过简单的处理，随机从各类中抽取20张组成测试集，训练集和测试集目录一样。文件夹组织如下：

2.pytorch加载图像数据集

2.1 torchvision.datasets.ImageFolder()读取图像

ImageFolder要求图像文件夹如有如下组织方式：

对应不同的类别会按照文件夹顺序生成0, 1, … , class - 1的标签，本数据集符合ImageFolder的要求，故采用此方式读取数据集。

ImageFolder类的定义如下：

class ImageFolder(root, transform=None, target_transform=None,

loader=default_loader, is_valid_file=None)

root: 包含类别文件夹的上一层目录，以本数据集的训练集为例，应该是D:\pokemon_dataset\train

transform: 对读出图像的数字矩阵进行的变换，输入为torchvision.transforms对象

target_transform: 对图像的类别（标签值）进行的变换

loader: 读取图像的方式，默认读取为RGB格式的PIL Image对象

is_valid_file:检查图像是否损坏

一般情况下我们只需使用root和transform两个参数，其他参数选择默认即可。接下来结合代码介绍ImageFolder的基本使用方法。

from torchvision import transforms

from torchvision.datasets import ImageFolder

from torch.utils.data import DataLoader

ROOT_TRAIN = r'D:\pokemon_dataset\train'

ROOT_TEST = r'D:\pokemon_dataset\test'

# 将图像RGB三个通道的像素值分别减去0.5再除以0.5，从而将所有像素值

# 固定到[-1.0, 1.0]范围内

normalize = transforms.Normalize([0.5, 0.5, 0.5], [0.5, 0.5, 0.5])

train_transform = transforms.Compose([

transforms.Resize((256, 256)), # 裁剪为256 * 256

transforms.RandomVerticalFlip(), # 随机垂直翻转

transforms.ToTensor(), # 将0-255范围的像素转为0-1.0范围的tensor

normalize])

test_transform = transforms.Compose([

transforms.Resize((256, 256)),

transforms.ToTensor(),

normalize])

train_dataset = ImageFolder(ROOT_TRAIN, transform=train_transform)

test_dataset = ImageFolder(ROOT_TEST, transform=test_transform)

因为打算测试多层神经网络的分类性能，所以这里对训练数据进行了resize、随机翻转的预处理，ToTensor的时候会自动将PIL Image对象转换为像素值在0-1.0范围内的(C, H, W)维张量(这里是(3, 256, 256))，所以在没有对图像预先进行像素统计的情况下，normalize可粗略地将像素值固定到[-1.0, 1.0]范围内，加快模型收敛。注意Compose的时候要传入转换流程的顺序列表。此外，对测试集进行的处理流程取决于选择的分类模型，这里去掉了随机翻转，其他和训练集保持一致。

print(len(train_dataset)) # 训练集大小（图像数量） out: 1067

print(train_dataset[0]) # 训练集第一张图像张量以及对应的标签，二维元组

print(train_dataset[0][0]) # 训练集第一张图像的张量

print(train_dataset[0][1]) # 训练集第一张图像的标签

ImageFolder返回的对象是一个包含数据集所有图像及对应标签构成的二维元组容器，支持索引和迭代，可作为torch.utils.data.DataLoader的输入。

测试过程中发现似乎不支持gif格式图像的读取，此外对于灰度图会默认以RGB三通道进行读取(C = 3)。

2.2 torch.utils.data.DataLoader()加载数据集

DataLoader的定义如下：

class torch.utils.data.DataLoader(dataset, batch_size=1, shuffle=False,

sampler=None, num_workers=0, collate_fn=,

pin_memory=False, drop_last=False)

这里只介绍几个比较重要的参数：

dataset: 要加载的数据集，可以传入ImageFolder的返回结果

batch_size: 每次读取图像的张数

shuffle: 是否打乱读取顺序

num_workers: 加载数据的子进程数

train_dataloader = DataLoader(train_dataset, batch_size=4, shuffle=True)

test_dataloader = DataLoader(test_dataset, batch_size=4, shuffle=False)

这里对训练数据和测试数据取batch_size = 4(据说batch_size取为2的整指数次方有好处)，训练集随机打乱，测试集不打乱，每个epoch都会把数据集所有图像都遍历一次。注意数据集如果不是batch_size的整数倍的话，最后一轮迭代的B会不足（小于）batch_size。

count = 0

for train_data in train_dataloader:

count += 1

imgs, target = train_data

print(imgs.size()) # out: (4, 3, 256, 256)

print(target, target.dtype) # out: tensor([4, 4, 3, 0]), torch.int64

if count == 1:

break

DataLoader返回一个可迭代对象，每次迭代的结果是一个二维元组，第一维是一个(B, C, H, W)维张量，保存的是图像数据；第二维是一个(B,)维张量，保存的是batch_size中对应图像的标签。

训练和测试时可对train_data的imgs和target进行相应的修改以适应网络和损失函数的要求。以多层神经网络为例，需要把imgs改成(B, C * H * W)维张量作为线性层的输入(torch.nn.Linear(C * H * W, h1_dim)；当损失函数选择torch.nn.CrossEntropyLoss()时，target不需要进行处理，因为交叉熵损失函数会自动先对target进行onehot编码，再计算对应的损失：

import torch as th

import numpy as np

out = th.randn((4, 5))

target = th.LongTensor([3, 1, 0, 1])

loss_f = th.nn.CrossEntropyLoss()

loss = loss_f(out, target)

out_np = out.numpy()

tar_np = np.array([[0, 0, 0, 1, 0],

[0, 1, 0, 0, 0],

[1, 0, 0, 0, 0],

[0, 1, 0, 0, 0]])

def softmax_cross_entropy(arr):

exp_arr = np.exp(arr)

for i in range(len(exp_arr)):

exp_arr[i] /= np.sum(exp_arr[i])

result = np.sum(np.log(exp_arr) * tar_np) / -4

return result

result = softmax_cross_entropy(out_np)

print(loss, result) # out: tensor(2.3829) 2.3828649520874023

原文链接:https://blog.csdn.net/caigi/article/details/105268908