使用torchvision下载数据并可视化

 

首先下载数据：

import torchvision.transforms
from torchvision import datasets
train_data = datasets.FashionMNIST(
    root="data",  # 保存的文件位置
    train=True,  # 是否选择训练数据，否的话表示测试数据
    download=True,  # 是否下载
    transform=torchvision.transforms.ToTensor(),  # 转换成Tensor
    target_transform=None   # 是否需要转换标签，None则使用原本的标签
)

test_data = datasets.FashionMNIST(
    root="data",
    train=False,  # 下载测试数据
    download=True,
    transform=torchvision.transforms.ToTensor(),
    target_transform=None
)

len(train_data),len(test_data)  # 查看数据长度
# 结果
# 60000, 10000
# 意味着有60000张训练图，10000张用于测试的图

image,label = train_data[0] # 选择第一张图作为待测图像
image,label   # 输出图像的Tensor和标签类别

运行的结果为：

(tensor([[[0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000,
           0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000,
           0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000,
           0.0000, 0.0000, 0.0000, 0.0000],……

         0.0000, 0.0000, 0.0000, 0.0000]]]),
 9)

从上方结果来看，训练集的第一张图的标签类别为9，标签是从0开始算的，所以是第10类

class_names = train_data.classes
class_names
# 结果为：
['T-shirt/top',
 'Trouser',
 'Pullover',
 'Dress',
 'Coat',
 'Sandal',
 'Shirt',
 'Sneaker',
 'Bag',
 'Ankle boot']

class_to_idx = train_data.class_to_idx
class_to_idx  # 返回一个不同标签及其对应索引的字典
# 结果为
{'T-shirt/top': 0,
 'Trouser': 1,
 'Pullover': 2,
 'Dress': 3,
 'Coat': 4,
 'Sandal': 5,
 'Shirt': 6,
 'Sneaker': 7,
 'Bag': 8,
 'Ankle boot': 9}

所以这张图的标签类别是'Ankle boot'，查看该图的大小：

# 检查图像的形状
image.shape,label
# 结果
#(torch.Size([1, 28, 28]), 9)
# 显然图1是一张黑白色的28*28的图像

下面，可视化这张图

import matplotlib.pyplot as plt
image, label = train_data[0]
print(f"Image shape:{image.shape}")
plt.imshow(image)

# 结果
# Image shape:torch.Size([1, 28, 28])

但是绘图报错：

 这是一个在机器学习中常见的错误，因为你的输入张量的形状与该张量的预期形状不符。

我们的数据格式首先是颜色通道1，然后是图像宽度28和高度28，但是matplotlib则希望先是图像的宽高然后再是颜色通道，因此需要改变形状再绘图。

对于灰度图像，我们可以使用.squeeze来消除这个额外的维度：

plt.imshow(image.squeeze())

运行结果为：

 我们再添加标签类别，使用cmap绘制灰度图，去除坐标轴：

 最后，我们绘制出更多的图像。我们使用随机性，看看数据集内的随机图像。

import torch
torch.manual_seed(42) # 设置随机种子
fig = plt.figure(figsize=(9,9))  # 创建一个大小为9x9的图像对象，用于存放后续绘制的图像网格
rows, cols = 4, 4   # 展示4*4=16张图
for i in range(1,rows*cols+1):   # 循环遍历每个网格的位置
    random_idx = torch.randint(0,len(train_data),size =[1]).item()
# 生成一个随机索引，用于从train_data数据集中随机选择一张图像和对应的标签
    # print(random_idx)
    img, label = train_data[random_idx] # 根据随机索引从train_data中取一张图像和对应的标签
    fig.add_subplot(rows, cols, i) # 在图像对象中添加一个子图，位置由rows和cols决定                                                    
    plt.imshow(img.squeeze(), cmap="gray")  # i是子图的索引
    plt.title(class_names[label])
    plt.axis(False);

我们看看最后的结果：

在后续的文章中，我们将介绍使用一个CV模型来对这些图进行分类。