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下的文件夹。保存结果如下图