PyTorch_torchvision---数据集下载\TensorBoard可视化

1. 引言

在深度学习入门阶段，数据集是绕不开的第一道坎。很多初学者卡在「数据从哪来、怎么加载、怎么查看」这一步，导致后续模型训练迟迟无法推进。

PyTorch 生态中有一个非常实用的工具库------torchvision，它内置了大量经典数据集（如 CIFAR-10、MNIST、Fashion-MNIST 等），只需几行代码就能自动下载并加载，配合 TensorBoard 还能直观地可视化数据内容。

本文将以 CIFAR-10 数据集 为例，带你走完「数据集介绍 → 下载使用 → 内容查看 → TensorBoard 可视化」的完整流程，帮助你把数据准备环节彻底打通。

2. torchvision 数据集介绍

2.1 什么是 torchvision

torchvision 是 PyTorch 官方提供的视觉工具库，包含：

常用数据集：CIFAR-10、MNIST、ImageNet 等；
图像变换工具 ：transforms 模块；
预训练模型：ResNet、VGG 等；
工具函数：图像读写、标注框处理等。

当我们写代码时，只需指定数据集名称和相应参数，torchvision 就会自动帮我们下载数据到本地，无需手动去官网找资源。

2.2 CIFAR-10 数据集简介

CIFAR-10 是一个非常经典的图像分类数据集，特点如下：

属性	说明
图片尺寸	32 × 32 彩色图片（RGB 三通道）
图片总数	60000 张
类别数量	10 类
训练集	50000 张
测试集	10000 张

10 个类别分别是：airplane（飞机）、automobile（汽车）、bird（鸟）、cat（猫）、deer（鹿）、dog（狗）、frog（青蛙）、horse（马）、ship（船）、truck（卡车）。

提示：32×32 的分辨率虽然不高，但足以用来学习图像分类的基本流程，且数据量适中，非常适合入门练习。

3. torchvision 数据集下载与使用

3.1 环境准备

在 Anaconda 终端中，先激活你的 Python 环境（本文以 py3.6.3 为例），再切换到项目所在路径：

bash 复制代码

conda activate py3.6.3
cd /path/to/your/project

3.2 下载 CIFAR-10 数据集

进入 Python 交互环境，导入 torchvision 并执行下载命令：

python 复制代码

import torchvision

train_set = torchvision.datasets.CIFAR10(
    root="./dataset",
    train=True,
    download=True
)

执行后，torchvision 会自动从官网下载 cifar-10-python.tar.gz 并解压到 ./dataset 目录下。如果数据集已经下载过，再次运行不会重复下载。

3.3 查看 CIFAR10 类的帮助文档

我们可以通过 help() 查看 CIFAR10 类的完整说明：

python 复制代码

import torchvision
help(torchvision.datasets.CIFAR10)

核心参数说明如下：

参数	类型	说明
`root`	string	数据集存放的根目录，若 `download=True` 且目录不存在，会自动创建并下载
`train`	bool	`True` 表示加载训练集，`False` 表示加载测试集
`transform`	callable	对每张图片应用的变换函数，如 `transforms.ToTensor()`
`target_transform`	callable	对标签（target）应用的变换函数
`download`	bool	是否从互联网下载数据集；已下载则不会重复下载

4. 查看 CIFAR-10 数据集内容

下载完成后，我们来实际查看数据集里的内容，验证数据是否加载正确。

4.1 加载训练集与测试集

python 复制代码

import torchvision

# root 为存放数据集的相对路径
train_set = torchvision.datasets.CIFAR10(
    root="./dataset",
    train=True,
    download=True
)

# train=True 是训练集，train=False 是测试集
test_set = torchvision.datasets.CIFAR10(
    root="./dataset",
    train=False,
    download=True
)

4.2 查看单张图片与标签

python 复制代码

print(test_set[0])        # 输出 (图片, 标签)，这里的 3 是 target
print(test_set.classes)   # 打印全部类别名称

img, target = test_set[0] # 分别取出图片和标签
print(img)                # 查看图片对象信息
print(target)             # 查看标签数值
print(test_set.classes[target])  # 3 号 target 对应的类别名称

img.show()                # 弹出窗口显示图片

运行结果示例：

text 复制代码

(<PIL.Image.Image image mode=RGB size=32x32 at 0x1A4275AAF28>, 3)
['airplane', 'automobile', 'bird', 'cat', 'deer', 'dog', 'frog', 'horse', 'ship', 'truck']
<PIL.Image.Image image mode=RGB size=32x32 at 0x1A4275AAA58>
3
cat

可以看到，test_set[0] 返回一个元组 (图片, 标签)，其中标签 3 对应类别 cat（猫）。

5. 使用 TensorBoard 可视化数据集

5.1 将图片转为 Tensor 并写入日志

TensorBoard 是 TensorFlow 推出的可视化工具，PyTorch 也提供了对应的接口 torch.utils.tensorboard.SummaryWriter，可以方便地查看图片、损失曲线、网络结构等。

python 复制代码

import torchvision
from torch.utils.tensorboard import SummaryWriter

# 定义数据变换：将 PIL 图片转为 Tensor 类型
dataset_transform = torchvision.transforms.Compose([
    torchvision.transforms.ToTensor()
])

# 将 ToTensor 应用到数据集中的每一张图片
train_set = torchvision.datasets.CIFAR10(
    root="./dataset",
    train=True,
    transform=dataset_transform,
    download=True
)

test_set = torchvision.datasets.CIFAR10(
    root="./dataset",
    train=False,
    transform=dataset_transform,
    download=True
)

writer = SummaryWriter("logs")

for i in range(10):
    img, target = test_set[i]
    writer.add_image("test_set", img, i)
    print(img.size())

writer.close()  # 一定要关闭 writer，否则图片可能无法显示

运行结果示例：

text 复制代码

torch.Size([3, 32, 32])
torch.Size([3, 32, 32])
torch.Size([3, 32, 32])
torch.Size([3, 32, 32])
torch.Size([3, 32, 32])
torch.Size([3, 32, 32])
torch.Size([3, 32, 32])
torch.Size([3, 32, 32])
torch.Size([3, 32, 32])
torch.Size([3, 32, 32])

可以看到，经过 ToTensor() 变换后，每张图片的尺寸变为 [3, 32, 32]，即 3 通道 × 32 宽 × 32 高，符合 CIFAR-10 的原始规格。

5.2 启动 TensorBoard 查看日志

在 Anaconda 终端中，激活环境后输入以下命令：

bash 复制代码

tensorboard --logdir=C:\Users\wangy\Desktop\03CV\logs

将终端输出的网址复制到浏览器地址栏并回车，即可在 TensorBoard 界面中查看刚才写入的图片日志。

注意：--logdir 路径要与代码中 SummaryWriter("logs") 生成的目录保持一致，否则看不到数据。

6. 总结

本文围绕 CIFAR-10 数据集，完整演示了 torchvision 的下载、加载、查看与可视化流程，核心要点如下：

torchvision 内置数据集：一行代码即可自动下载，省去手动找数据的麻烦；
CIFAR-10 基本属性：60000 张 32×32 彩色图、10 个类别、训练/测试各 50000/10000 张；
数据查看 ：通过索引访问 (img, target) 元组，用 classes 列表映射类别名称；
TensorBoard 可视化 ：借助 SummaryWriter 将图片写入日志，浏览器中直观查看。

掌握了这套流程，后续无论是训练自己的模型，还是复现经典论文，数据准备环节都能轻松搞定。接下来，就可以放心地进入模型搭建与训练阶段了。