数据集读取与划分，ImageFolder()，自定义数据集，TensorDataset，StratifiedShuffleSplit

目录

导包

数据集

下载数据集

数据集特点分析

torchvision.datasets.ImageFolder()

数据集整理

思路

根据图片名读标签

建立标签子文件夹

数据集划分

调用数据集处理函数

读取数据集 

torchvision.datasets.ImageFolder()源码及解读

源码

解读

torchvision.datasets.ImageFolder()的特点

手写ImageFolder()

自定义数据集

数据集处理

将标签转换成为int类型

自定义数据集函数 

读取数据集 

数据集划分函数

train_test_split函数

用法

读取数据集

存在的问题

StratifiedShuffleSplit函数

用法

读取数据集

为什么要进行reset_index()操作

探析

random_split()

用法

划分

读取数据集

探析

其他dataset

TensorDataset

源码

 建立数据集

使用dataloader

一个问题：对单个tensor进行包装

导包

import torch
import torchvision
import torchvision.transforms as transforms
import torch.nn as nn
import pandas as pd
import os
import collections
import shutil
import math
from torch.utils.data import DataLoader, Dataset
from PIL import Image


data_dir = 'data\dog-breed-identification'  # 数据集所在文件夹
label_csv = 'labels.csv'  # 标签文件夹名

数据集

原教程网站：13.13. 实战 Kaggle 比赛：图像分类 (CIFAR-10) — 动手学深度学习 2.0.0-beta1 documentation13.14. 实战Kaggle比赛：狗的品种识别（ImageNet Dogs） — 动手学深度学习 2.0.0-beta1 documentation13.13. 实战 Kaggle 比赛：图像分类 (CIFAR-10) — 动手学深度学习 2.0.0-beta1 documentation

参考：动手学深度学习Kaggle：图像分类 (CIFAR-10和Dog Breed Identification)_iwill323的博客-CSDN博客

下载数据集

数据集网址是CIFAR-10 - Object Recognition in Images | KaggleDog Breed Identification | KaggleCIFAR-10 - Object Recognition in Images | Kaggle

下载数据集，在../data中解压下载的文件后，你将在以下路径中找到整个数据集：

•     ../data/dog-breed-identification/labels.csv
•     ../data/dog-breed-identification/sample_submission.csv
•     ../data/dog-breed-identification/train
•     ../data/dog-breed-identification/test

文件夹train/和test/分别包含训练和测试狗图像，labels.csv包含训练图像的标签，其中train文件夹含有样本图片的如下图，图像文件的名称是杂乱的

数据集特点分析

比赛数据集分为训练集和测试集，分别包含RGB（彩色）通道的10222张、10357张JPEG图像。 在训练数据集中，有120种犬类，如拉布拉多、贵宾、腊肠、萨摩耶、哈士奇、吉娃娃和约克夏等。​

• 用pandas读取trainLabels.csv文件

df = pd.read_csv(os.path.join(data_dir, label_csv))
df.head()

• 标签个数

breeds = df.breed.unique()
len(breeds)

120

• 训练集中每个种类的样本有多少

count_train = collections.Counter(df['breed'])
count_train.most_common() 

[('scottish_deerhound', 126),
 ('maltese_dog', 117),
 ('afghan_hound', 116),
 ……
 ('komondor', 67),
 ('brabancon_griffon', 67),
 ('eskimo_dog', 66),
 ('briard', 66)]

torchvision.datasets.ImageFolder()

数据集整理

思路

torchvision.datasets.ImageFolder()要求根目录下建立分类标签子文件夹，每个子文件夹下归档对应标签的图片，因此需要给每个标签建立文件夹，并且遍历样本，将每个样本复制到对应的文件夹中。本例在归档图片的时候，顺便把数据集划分了

根据图片名读标签

为了在根目录下按类别建立子文件夹，需要在读取每个样本图片名的时候，获得对应的类别标签label。然而，pandas一般根据表的index或者行数来选择数据，我没找到根据某一列的值索引其他列的数据的方法。教程根据一列的数据索引另一列的数据，下面的read_csv_labels()函数起到这样的作用，read_csv_labels函数返回的是一个字典格式的变量，该变量根据name可以索引label。

def read_csv_labels(fname):
    """读取fname来给标签字典返回一个文件名"""
    with open(fname, 'r') as f:
        # 跳过文件头行(列名)
        lines = f.readlines()[1:]
    tokens = [l.rstrip().split(',') for l in lines]
    return dict(((name, label) for name, label in tokens))

建立标签子文件夹

copyfile函数将图片从原位置filename复制到对应文件夹之下，只需要将target_dir指定为标签文件夹名字就行。

def copyfile(filename, target_dir):
    """将文件复制到目标目录"""
    os.makedirs(target_dir, exist_ok=True)  # 文件夹不存在则创建
    shutil.copy(filename, target_dir)

数据集划分

数据集只含有train和test数据集，而我们在训练的时候，一般还包含验证集，所以要划分出验证集处理。使用Google Colab这样的平台时，我们经常会将训练集、测试集、验证集压缩并上传，所以有时候要将它们划分、保存在不同的文件夹。

• 定义reorg_train_valid函数来将验证集从原始的训练集中拆分出来。 此函数中的参数valid_ratio是验证集中的样本数与原始训练集中的样本数之比。 更具体地说，令n等于样本最少的类别中的图像数量，而r是比率。 验证集将为每个类别拆分出max(⌊nr⌋,1)张图像。以valid_ratio=0.1为例，由于原始的训练集有50000张图像，因此train_valid_test/train路径中将有45000张图像用于训练，而剩下5000张图像将作为路径train_valid_test/valid中的验证集。
• 定义reorg_test函数将测试集数据复制到新文件夹，注意test文件夹下面也要有一个子文件夹（unknown）作为分类文件夹，否则torchvision.datasets.ImageFolder()会报错。因为ImageFolder()的find_classes()函数要从根文件夹下读取文件夹的名称，生成类别列表，没有这个列表就会导致错误

def reorg_train_valid(data_dir, labels, valid_ratio):
    """将验证集从原始的训练集中拆分出来"""
    # 训练数据集中样本最少的类别中的样本数
    n = collections.Counter(labels.values()).most_common()[-1][1]
    # 验证集中每个类别的样本数
    n_valid_per_label = max(1, math.floor(n * valid_ratio))
    label_count = {}
    for train_file in os.listdir(os.path.join(data_dir, 'train')):
        label = labels[train_file.split('.')[0]] # 根据文件名索引label
        fname = os.path.join(data_dir, 'train', train_file)
        copyfile(fname, os.path.join(data_dir, 'train_valid_test',
                                     'train_valid', label))
        if label not in label_count or label_count[label] < n_valid_per_label:
            copyfile(fname, os.path.join(data_dir, 'train_valid_test',
                                         'valid', label))
            label_count[label] = label_count.get(label, 0) + 1
        else:
            copyfile(fname, os.path.join(data_dir, 'train_valid_test',
                                         'train', label))
    return n_valid_per_label


def reorg_test(data_dir):
    """在预测期间整理测试集，以方便读取"""
    for test_file in os.listdir(os.path.join(data_dir, 'test')):
        copyfile(os.path.join(data_dir, 'test', test_file),
                 os.path.join(data_dir, 'train_valid_test', 'test',
                              'unknown'))

调用数据集处理函数

labels.values()的格式是<class 'builtin_function_or_method'>，可以用于collections.Counter()方法 

def reorg_cifar10_data(data_dir, label_csv, valid_ratio):
    labels = read_csv_labels(os.path.join(data_dir, label_csv))
    reorg_train_valid(data_dir, labels, valid_ratio)
    reorg_test(data_dir)


batch_size = 128
valid_ratio = 0.1
reorg_cifar10_data(data_dir, label_csv, valid_ratio)

代码执行的效果是，创建了四个文件夹，分别是test，train（9502个样本），valid（720个样本）和train_valid，其中train_valid是train和valid的合集。建立train_valid文件夹是因为，使用验证集筛选出最佳超参数之后，再使用train_valid训练一遍，得到最终模型

每一个文件夹下按照类别创建了120个分类文件夹，这是torchvision.datasets.ImageFolder()函数的要求。

读取数据集 

读取由原始图像组成的数据集，每个样本都包括一张图片和一个标签。注意，当验证集在超参数调整过程中用于模型评估时，不应引入图像增广的随机性，所以valid数据集使用的transform是transform_test

train_ds, train_valid_ds = [torchvision.datasets.ImageFolder(
    os.path.join(data_dir, 'train_valid_test', folder),
    transform=train_transform) for folder in ['train', 'train_valid']]

valid_ds, test_ds = [torchvision.datasets.ImageFolder(
    os.path.join(data_dir, 'train_valid_test', folder),
    transform=test_transform) for folder in ['valid', 'test']]

train_iter, train_valid_iter = [torch.utils.data.DataLoader(
    dataset, batch_size, shuffle=True, drop_last=True)
    for dataset in (train_ds, train_valid_ds)]

valid_iter = torch.utils.data.DataLoader(valid_ds, batch_size, shuffle=False,
                                         drop_last=True)

test_iter = torch.utils.data.DataLoader(test_ds, batch_size, shuffle=False,
                                        drop_last=False)

其中用到的图像增广：

img_size = 224  # 也可以是其他值
train_transform = transforms.Compose([    
    transforms.RandomResizedCrop(img_size, ratio=(3.0/4.0, 4.0/3.0)),
    transforms.RandomHorizontalFlip(),
    transforms.RandomRotation(30),
    transforms.ColorJitter(brightness=0.4, contrast=0.4, saturation=0.4),
    transforms.ToTensor(),
    transforms.Nor