PyTorch入门学习（五）：torchvision中dataset的使用（以CIFAR10和MNIST为例）

目录

一、引言

二、CIFAR10数据集

三、测试数据集和数据展示

---

一、引言

在深度学习项目中，数据集的准备和处理是一项繁重的工作。torchvision是PyTorch中的一个模块，提供了多种流行的数据集，方便用户加载和处理数据。本文将以CIFAR10和MNIST数据集为例，演示如何使用torchvision中的数据集。

二、CIFAR10数据集

CIFAR10是一个包含10个类别的图像分类数据集，每个类别包含6000张32x32像素的彩色图像。以下是如何使用torchvision加载CIFAR10数据集的代码：

import torchvision
from torch.utils.tensorboard import SummaryWriter
from torch.utils.data import DataLoader

# 准备测试数据集
test_data = torchvision.datasets.CIFAR10(root="D:\\Python_Project\\pytorch/dataset2", train=False, transform=torchvision.transforms.ToTensor())

test_loader = DataLoader(dataset=test_data, batch_size=64, shuffle=True, num_workers=0, drop_last=True)

• torchvision.datasets.CIFAR10用于加载CIFAR10数据集，其中root参数指定了数据集的根目录，train参数设置为False表示加载测试数据集，transform参数将图像转换为张量格式。
• DataLoader用于将数据集分成小批次，以便于模型训练。在示例中，设置了每批包含64张图像，启用了数据的随机顺序洗牌，并设置num_workers为0以避免并行处理。

三、测试数据集和数据展示

接下来，从CIFAR10数据集中获取一张图像和其对应的目标标签，以及使用TensorBoard进行数据可视化：

# 测试数据集中第一张图片及其target
img, target = test_data[0]
print(img.shape)
print(target)

writer = SummaryWriter("logs")
for epoch in range(2):
    step = 0
    for data in test_loader:
        imgs, targets = data
        writer.add_images("Epoch:{}".format(epoch), imgs, step)
        step = step + 1
writer.close()

• 从test_data中获取第一张图像和其对应的目标标签。
• 使用SummaryWriter创建一个TensorBoard写入器，以便于数据可视化。
• 通过循环迭代数据批次，将图像数据添加到TensorBoard中，以不同的命名空间标识不同的迭代步骤。

完整代码如下：

import torchvision
from torch.utils.tensorboard import SummaryWriter
from torch.utils.data import DataLoader

# 准备测试数据集
# 使用torchvision加载CIFAR10测试数据集，将数据集保存在指定的根目录，并将图像转换为张量形式
test_data = torchvision.datasets.CIFAR10(root="D:\\Python_Project\\pytorch/dataset2", train=False, transform=torchvision.transforms.ToTensor())

# 创建数据加载器
test_loader = DataLoader(dataset=test_data, batch_size=64, shuffle=True, num_workers=0, drop_last=True)

# 测试数据集中第一张图片及其target
img, target = test_data[0]
print(img.shape)
print(target)

# 创建一个TensorBoard写入器，将日志保存到名为"logs"的文件夹
writer = SummaryWriter("logs")

# 迭代两个周期（epochs）
for epoch in range(2):
    step = 0
    # 遍历测试数据加载器
    for data in test_loader:
        imgs, targets = data
        # 添加图像数据到TensorBoard，使用不同的命名空间标识不同的迭代步骤
        writer.add_images("Epoch:{}".format(epoch), imgs, step)
        step = step + 1

# 关闭TensorBoard写入器
writer.close()

参考资料：

视频教程：PyTorch深度学习快速入门教程（绝对通俗易懂！）【小土堆】