CIFAR-10数据集共有60 000幅彩色图像，这些图像是32×32像素的，分为10类，每类6 000幅图，如图7-9所示。这里面有50 000幅图用于训练，构成了5个训练批，每一批10 000幅图；另外，10 000幅用于测试，单独构成一批。测试批的数据取自100类中的每一类，每一类随机取1000幅。抽剩下的就随机排列组成训练批。注意，一个训练批中的各类图像的数量并不一定相同，总的来看，训练批每一类都有5 000幅图。

图7-9  CIFAR-10数据集

读者自行搜索CIFAR-10数据集下载地址，进入下载页面后，选择下载方式，如图7-10所示。

图7-10  下载方式

由于PyTorch 2.0采用Python语言编程，因此选择Python Version的版本下载。下载之后解压缩，得到如图7-11所示的文件。

图7-11  得到的文件

data_batch_1~data_batch_5是划分好的训练数据，每个文件中包含10 000幅图片，test_batch是测试集数据，也包含10 000幅图片。

读取数据的代码如下：

import pickle

def load_file(filename):

    with open(filename, 'rb') as fo:

        data = pickle.load(fo, encoding='latin1')

    return data

首先定义读取数据的函数，这几个文件都是通过 pickle 产生的，所以在读取的时候也要用到这个包。返回的data是一个字典，先来看这个字典里面有哪些键。

data = load_file('data_batch_1')

print(data.keys())

输出结果如下：

dict_key3(['batch_label', 'labels', 'data', 'filenames'])

具体说明如下。

1. batch_label：对应的值是一个字符串，用来表明当前文件的一些基本信息。
2. labels：对应的值是一个长度为10 000的列表，每个数字取值范围为0~9，代表当前图片所属的类别。
3. data：10000×3072的二维数组，每一行代表一幅图片的像素值。
4. filenames：长度为10 000的列表，里面每一项是代表图片文件名的字符串。

完整的数据读取函数如下：

【程序7-1】

i

import pickle
import numpy as np
import os

def get_cifar10_train_data_and_label(root=""):
    def load_file(filename):
        with open(filename, 'rb') as fo:
            data = pickle.load(fo, encoding='latin1')
        return data

    data_batch_1 = load_file(os.path.join(root, 'data_batch_1'))
    data_batch_2 = load_file(os.path.join(root, 'data_batch_2'))
    data_batch_3 = load_file(os.path.join(root, 'data_batch_3'))
    data_batch_4 = load_file(os.path.join(root, 'data_batch_4'))
    data_batch_5 = load_file(os.path.join(root, 'data_batch_5'))
    dataset = []
    labelset = []
    for data in [data_batch_1, data_batch_2, data_batch_3, data_batch_4, data_batch_5]:
        img_data = (data["data"])
        img_label = (data["labels"])
        dataset.append(img_data)
        labelset.append(img_label)
    dataset = np.concatenate(dataset)
    labelset = np.concatenate(labelset)
    return dataset, labelset

def get_cifar10_test_data_and_label(root=""):
    def load_file(filename):
        with open(filename, 'rb') as fo:
            data = pickle.load(fo, encoding='latin1')
        return data

    data_batch_1 = load_file(os.path.join(root, 'test_batch'))
    dataset = []
    labelset = []
    for data in [data_batch_1]:
        img_data = (data["data"])
        img_label = (data["labels"])
        dataset.append(img_data)
        labelset.append(img_label)
    dataset = np.concatenate(dataset)
    labelset = np.concatenate(labelset)
    return dataset, labelset

def get_CIFAR10_dataset(root=""):
    train_dataset, label_dataset = get_cifar10_train_data_and_label(root=root)
    test_dataset, test_label_dataset = get_cifar10_train_data_and_label(root=root)
    return train_dataset, label_dataset, test_dataset, test_label_dataset

if __name__ == "__main__":
    train_dataset, label_dataset, test_dataset, test_label_dataset = get_CIFAR10_dataset(root="../dataset/cifar-10-batches-py/")

    train_dataset = np.reshape(train_dataset,[len(train_dataset),3,32,32]). astype(np.float32)/255.
    test_dataset = np.reshape(test_dataset,[len(test_dataset),3,32,32]). astype(np.float32)/255.
    label_dataset = np.array(label_dataset)
    test_label_dataset = np.array(test_label_dataset)

其中的root是下载数据解压后的目录参数，os.join函数将其组合成数据文件的位置。最终返回训练文件和测试文件以及它们对应的label。需要说明的是，提取出的文件数据格式为[-1,3072]，因此需要重新对数据维度进行调整，使之适用于模型的输入。