Pytorch读取MNIST数据集并保存里面的图片

绝地灬酷狼 2022-09-03 05:21 165阅读 0赞

# Pytorch读取MNIST数据集并保存里面的图片 #

最近在学习Pytorch编程，在里面训练模型的时候需要用到比较常用的MNIST数据集，在此做出总结，以便后面不时之需。

## 1.导入相关库 ##

import torch
    import torch.nn as nn
    from torchvision import datasets, transforms
    import numpy as np
    import matplotlib.pyplot as plt

## 2.下载MNIST训练数据和测试数据 ##

"""加载Mnist手写数据集训练集"""
    mnist_train = datasets.MNIST('./data_mnist',
                       train=True,
                       download=True,
                       transform=transforms.Compose([transforms.ToTensor()]))
    """加载Mnist手写数据集测试集合"""
    mnist_test = datasets.MNIST('./data_mnist',
                       train=False,
                       download=True,
                       transform=transforms.Compose([transforms.ToTensor()]))

注意上面两个代码，train=True时候是训练集，为False的时候是测试集，最后的transform的参数是将数据变成Tensor格式的数据。

## 3.看一下训练集和测试集的大小 ##

train_data = torch.utils.data.DataLoader(mnist_train, batch_size=1, shuffle=True)
    test_data = torch.utils.data.DataLoader(mnist_test, batch_size=1, shuffle=True)
    print(len(train_data))
    print(len(test_data))

将数据加载到变量train\_data和test\_data中，用的是torch自带的函数，生成类似一个序列的格式，shuffle=True是打乱顺序，输出结果分别为60000， 10000。

## 4.保存数据为图片格式 ##

# 先创建一个文件夹
    import os
    os.mkdir('images')
    # 使用plt.imsave()进行保存，本次以保存测试数据为例
    for i in range(len(test_data)):
        data, target = next(iter(test_data)) # 迭代器
        new_data = data[0][0].clone().numpy() # 拷贝数据
        plt.imsave('images/'+str(i)+str(target)+'.png', new_data)
        print(target)

运行此代码就可以把图片保存到images下的文件夹。保存结果如下图  
![保存的测试集][watermark_type_ZmFuZ3poZW5naGVpdGk_shadow_10_text_aHR0cHM6Ly9ibG9nLmNzZG4ubmV0L3NpbmF0XzQzMDA5OTgy_size_16_color_FFFFFF_t_70]

[watermark_type_ZmFuZ3poZW5naGVpdGk_shadow_10_text_aHR0cHM6Ly9ibG9nLmNzZG4ubmV0L3NpbmF0XzQzMDA5OTgy_size_16_color_FFFFFF_t_70]: /images/20220829/eec8d3f5457f4f668bb05fd20ba0e0e2.png