pytorch的dataloader函数

Nina_717

已于 2025-02-27 18:44:37 修改

阅读量1.3k

点赞数 26

分类专栏： pytorch入门文章标签： pytorch 人工智能 python

于 2025-02-27 16:26:54 首次发布

本文链接：https://blog.youkuaiyun.com/qq_62585870/article/details/145907448

版权

pytorch入门专栏收录该内容

6 篇文章

订阅专栏

例子代码使用了 PyTorch 和 PyTorch Vision 来加载 CIFAR-10 数据集，并将数据集中的图像可视化到 TensorBoard 中。以下是对代码的详细解释：

1. 导入必要的模块

Python复制

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

torchvision: PyTorch 的扩展库，提供了一系列工具，包括常用的计算机视觉数据集（如 CIFAR-10）和数据转换操作（如 ToTensor）。
DataLoader: PyTorch 的数据加载工具，用于将数据集分割成小批量（batch），并进行数据预处理（如打乱数据、多线程加载等）。
SummaryWriter: TensorBoard 的写入工具，用于记录训练过程中的各种信息（如图像、标量、直方图等）。

2. 加载 CIFAR-10 数据集

Python复制

test_dataset = torchvision.datasets.CIFAR10("./data", train=False, transform=torchvision.transforms.ToTensor())

CIFAR10: CIFAR-10 是一个经典的计算机视觉数据集，包含 10 个类别的图像（如飞机、汽车、猫、狗等）。
train=False: 表示加载的是测试集（False）而不是训练集（True）。
transform=torchvision.transforms.ToTensor(): 将图像数据从 PIL 图像格式（常见的图像格式）转换为 PyTorch 的张量格式。张量格式通常更适合深度学习模型的处理，其形状为 (C, H, W)（通道数，高度，宽度）。

3. 创建 DataLoader

Python复制

test_loader = DataLoader(dataset=test_dataset, batch_size=4, shuffle=False, num_workers=0, drop_last=False)

dataset=test_dataset: 指定加载的数据集。
batch_size=4: 每次加载的图像数量为 4 张（批量大小）。
shuffle=False: 是否在每个 epoch（遍历数据集的次数）中随机打乱数据。这里设为 False，表示不打乱数据。
num_workers=0: 数据加载的线程数。设为 0 表示不使用额外的线程。
drop_last=False: 如果数据集的大小不能被批量大小整除，是否丢弃最后一个较小的批量。这里设为 False，表示保留最后一个较小的批量。

4. 初始化 SummaryWriter

Python复制

writer = SummaryWriter("dataloader")

SummaryWriter: 初始化 TensorBoard 的写入工具，指定日志文件的保存路径为 "dataloader"。
日志路径: TensorBoard 的日志文件会保存在当前目录下名为 dataloader 的文件夹中。

5. 遍历数据并写入 TensorBoard

Python复制

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

for epoch in range(2): 外层循环，表示遍历数据集 2 次（模拟 2 个 epoch）。
for data in test_loader: 内层循环，逐批加载数据。
imgs, targets = data: 将每批数据拆分为图像（imgs）和标签（targets）。
writer.add_images: 将图像数据写入 TensorBoard。参数说明：
- "Epoch:{}".format(epoch): 记录的标签名称，会随 epoch 动态生成。
- imgs: 图像张量，形状为 (batch_size, C, H, W)。
- step: 记录的步数（每个 batch 对应一个步数）。
step += 1: 每处理一个 batch，步数递增。

6. 关闭 SummaryWriter

Python复制

writer.close()

关闭 SummaryWriter，确保所有数据都被正确写入到 TensorBoard 的日志文件中。

总结

这段代码的主要目的是加载 CIFAR-10 数据集，并将数据集中的图像以批量的形式可视化到 TensorBoard 中。
通过 DataLoader 和 SummaryWriter，可以方便地加载和可视化数据。
TensorBoard 的数据可以使用以下命令查看：

bash复制
```
tensorboard --logdir dataloader
```
打开浏览器并访问 http://localhost:6006 即可查看可视化结果。