跳至主要內容

二、PyTorch—DataLoader & Transforms

原创Xenny约 1114 字大约 5 分钟机器学习机器学习PyTorch

二、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__,接收索引并返回一个元组(x,y)(x,y),即数据样本和标签值。而__len__返回这个数据集的长度。

  • 对于一些常见的数据类型,我们一般不需要自己编写逻辑,在torch.utils.data中包含了多个以及实现好了的数据集包装类,例如上一篇博文中的TensorDataset

  • 调用逻辑
    调用逻辑

    完整的逻辑可以看上图,可以按照 pytorch笔记三:pytorch数据读取机制(DataLoader)与图像预处理模块(transforms)open in new window 这篇文章自己动调一下就清楚了。

    简单来说DataLoader中的Sampler和shuffle定义了每次读取哪一个索引和一个批次读取多少,Dataset定义了第i个数据如何读取出来。

Transforms

  • 这是PyTorch中用于图像预处理和数据增强的模块,包含在torchvision子包中。transforms中的函数可以分为两大类

    1. 对PIL图像对象的操作函数,剪裁、旋转、缩放等等;
    2. 对张量对象的操作函数。
  • 这里我们依然用一个例子来说明数据读取和预处理工作。

    以下是一个使用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 classificationopen in new window,其中我们使用

    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%的正确率。