二、PyTorch—DataLoader & Transforms
二、PyTorch—DataLoader & Transforms
- 在训练时,我们读取各种不同类型的数据以及需要对数据做一些预处理,这里介绍的便是PyTorch中的数据读取和图像预处理模块。
DataLoader
在上一篇博文文末例子中我们已经使用了
torch.utils.data.DataLoader
来处理数据,这个类的初始化参数为DataLoader(dataset, batch_size=1, shuffle=False, sampler=None, num_workers=0, collate_fn=<function default_collate>, pin_memory=False, drop_last=False)
其中我们平时会用到的参数如下
dataset
(Dataset) – 加载数据的数据集;batch_size
(int, optional) – 每个batch加载多少个样本(默认: 1);shuffle
(bool, optional) – 设置为True时会在每个epoch重新打乱数据(默认: False);num_workers
(int, optional) – 用多少个子进程加载数据。0表示数据将在主进程中加载(默认: 0);drop_last
(bool, optional) – 如果数据集大小不能被batch size整除,则设置为True后可删除最后一个不完整的batch。如果设为False并且数据集的大小不能被batch size整除,则最后一个batch将更小。(默认: False)。
Dataset
在
DataLoader
中规定了数据从哪里来,以及数据的打乱方式,每个批次多少数据等,而具体的数据又需要一个新的封装类来规定每一条数据的具体读取方式,即torch.utils.data.Dataset
类。一个Dataset类需要实现两个方法
class custom_dataset(Dataset): def __getitem__(self, index): pass def __len__(self): pass
对于
__getitem__
,接收索引并返回一个元组,即数据样本和标签值。而__len__
返回这个数据集的长度。对于一些常见的数据类型,我们一般不需要自己编写逻辑,在
torch.utils.data
中包含了多个以及实现好了的数据集包装类,例如上一篇博文中的TensorDataset
。完整的逻辑可以看上图,可以按照 pytorch笔记三:pytorch数据读取机制(DataLoader)与图像预处理模块(transforms) 这篇文章自己动调一下就清楚了。
简单来说DataLoader中的Sampler和shuffle定义了每次读取哪一个索引和一个批次读取多少,Dataset定义了第i个数据如何读取出来。
Transforms
这是PyTorch中用于图像预处理和数据增强的模块,包含在
torchvision
子包中。transforms
中的函数可以分为两大类- 对PIL图像对象的操作函数,剪裁、旋转、缩放等等;
- 对张量对象的操作函数。
这里我们依然用一个例子来说明数据读取和预处理工作。
以下是一个使用ResNet18处理猫狗图片二分类问题的代码,其中ResNet18是一个经典的卷积神经网络。
import os import torch import torch.nn as nn from torchvision.models import resnet18 import torchvision.transforms as transforms from torch.utils.data import Dataset, DataLoader import PIL.Image class CustomDataset(Dataset): def __init__(self, path, transform) -> None: super().__init__() self.transform = transform self.imgs = self.make_dataset(path) def __getitem__(self, index): img_path, label = self.imgs[index] img = PIL.Image.open(img_path).convert('RGB') img = self.transform(img) return img, label def __len__(self): return len(self.imgs) @staticmethod def make_dataset(path): datasets = [] class_to_idx = {'cats': 0, 'dogs': 1} for root, dirs, _ in os.walk(path): for sub_dir in dirs: imgs = list(os.listdir(os.path.join(root, sub_dir))) for img in imgs: datasets.append(( os.path.join(root, sub_dir, img), class_to_idx[sub_dir] )) return datasets mean = [0.5875, 0.5591, 0.4995] std = [0.2898, 0.2873, 0.3088] train_transform = transforms.Compose([transforms.Resize((255,255)), transforms.RandomCrop(255, padding=4), transforms.ToTensor(), transforms.Normalize(mean, std) ]) valid_transform = transforms.Compose([transforms.Resize((255,255)), transforms.ToTensor(), transforms.Normalize(mean, std) ]) train_dataset = CustomDataset('./datasets/train', transform=train_transform) valid_dataset = CustomDataset('./datasets/test', transform=valid_transform) train_loader = DataLoader(dataset=train_dataset, batch_size=4, shuffle=True) valid_loader = DataLoader(dataset=valid_dataset, batch_size=4, shuffle=True) model = resnet18(weights=False, num_classes=2).to(device="cuda") criterion = nn.CrossEntropyLoss() optimizer = torch.optim.SGD(model.parameters(), lr=0.0001) for epoch in range(100): size = len(train_loader.dataset) for batch, (X, y) in enumerate(train_loader): X, y = X.cuda(), y.cuda() pred = model(X) loss = criterion(pred, y) optimizer.zero_grad() loss.backward() optimizer.step() if epoch % 10 == 0: model.eval() correct = 0 size = len(valid_loader.dataset) with torch.no_grad(): for X, y in valid_loader: X, y = X.cuda(), y.cuda() pred = model(X) correct += (pred.argmax(1) == y).type(torch.float).sum().item() print(f'Epoch {epoch}/100 Acc = {correct/size}')
数据集来自Cats and Dogs image classification,其中我们使用
train_transform = transforms.Compose([transforms.Resize((255,255)), transforms.RandomCrop(255, padding=4), transforms.ToTensor(), transforms.Normalize(mean, std) ])
创建了一个transform处理链,分别是
transforms.Resize 改变图像大小 transforms.RandomCrop 对图像进行像素填充再随机裁剪,目的是增加泛化能力 transforms.ToTensor 将图像转换成张量,并对像素进行归一化 transforms.Normalize 将数据进行标准化
在我们自定义的
CustomDataset
类中,每次取数据时即用PIL打开图片,再使用这个处理链转化为张量,对于数据集没有进行剪裁操作。在本地上述模型达到63%的正确率。