pytorch笔记02-数据读取

pytorch不仅内置的一些数据集方便进行实验，并且还提供了像torchvision.datasets.ImageFolder(),torch.utils.data.Dataset()，torch.utils.data.DataLoader()等函数便于我们读取自己的数据集。
当数据如以下结构时，ImageFolder非常方便的为我们处理分类问题：

PetImages/dog/000.jpg
PetImages/dog/001.jpg
...
PetImages/cat/000.jpg
PetImages/cat/001.jpg
PetImages/cat/002.jpg

这时只需要将PetImages文件夹的路径输入ImageFolder函数即可。
但是，例如猫狗大战数据集，只有训练集，没有验证集。这样就需要我们将数据集拆分成训练集和验证集了，此时ImageFolder就无能为力了。此时需要自己重写torch.utils.data.Dataset()类，实现训练集和测试集的拆分。
代码如下：

import os
from PIL import Image
from torch.utils import data
import torchvision
import numpy as np
from torchvision import transforms as T
import random
import matplotlib.pyplot as plt


class DogCat(data.Dataset):

    def __init__(self, root, transfrom=None, train=True):
        self.train = train
        class_dirs = [os.path.join(root, i) for i in os.listdir(root)]
        imgs = []
        # 将所有类别的图片存到列表里
        for i in class_dirs:
            imgs += [os.path.join(i, img) for img in os.listdir(i)]
        # 将数据打乱，以便分割
        random.shuffle(imgs)
        imgs_mun = len(imgs)
        # 训练集：验证集 = 7 ：3
        if self.train:
            self.imgs = imgs[:int(0.7*imgs_mun)]
        else:
            self.imgs = imgs[int(0.7*imgs_mun):]
        # print(self.imgs[0], len(self.imgs))
        # 默认的数据转换操作
        if transfrom is None:
            normalize = T.Normalize(mean=[0.5, 0.5, 0.5],
                                    std=[0.5, 0.5, 0.5])
            if self.train:
                self.transforms = T.Compose([
                    T.Resize(256),
                    T.RandomResizedCrop(224),
                    T.RandomHorizontalFlip(),
                    T.ToTensor(),
                    normalize
                ])
            else:
                self.transforms = T.Compose([
                    T.Resize(224),
                    T.CenterCrop(224),
                    T.ToTensor(),
                    normalize
                ])

    def __getitem__(self, index):
        img_path = self.imgs[index]
        label = 1 if img_path.split('\\')[-2] == 'Dog' else 0
        # print(img_path.split('\\')[-2])
        data = Image.open(img_path)
        data = self.transforms(data)
        return data, label

    def __len__(self):
        return len(self.imgs)

# 定义一个显示图片的函数，对划分的数据集进行测试
def imshow(img):
    img = img / 2 + 0.5
    npimg = img.numpy()
    plt.imshow(np.transpose(npimg, (1, 2, 0)))
    plt.show()

if __name__ == '__main__':

    root = '.\\PetImages'
    # 加载训练集
    train_dataset = DogCat(root, train=True)
    # 加载验证集
    val_dataset = DogCat(root, train=False)
    # 训练集的batch loader
    train_loader = data.DataLoader(train_dataset,
                                   batch_size=4,
                                   shuffle=True,
                                   num_workers=4)
    # 测试数据划分是否正确
    classes = ['cat', 'dog']
    imgs, labels = iter(train_loader).next()
    imshow(torchvision.utils.make_grid(imgs))
    print(' '.join('%5s' % classes[labels[j]] for j in range(4)))