目录
一、引言
在深度学习项目中,数据集的准备和处理是一项繁重的工作。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()
参考资料: