用pytorch加载自己的数据集（使用Dataloader替代Torchvision.datasets.MNIST）（附项目源码）

参考：https://blog.csdn.net/winycg/article/details/90318371

• 省流：

当使用dataloader函数切割数据集与标签时，可以以 列表[元组(tensor数据,int标签)]的格式输入。

例：

将文件夹中的1000张图片输入代码，每张图片尺寸1*28*28，

标签以int格式输入。

分割批次大小：128

# 数据    :torch.Size([1, 28, 28])   | 图像
# 标签    :int                       | 图像对应的类型标签

[(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, 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]]]), 6),        # 末尾数字为图像的标签

...
(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, 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]]]), 7)]        # 末尾数字为图像的标签

后输入dataloader()函数，得到被分割的数据集：

trainloader = torch.utils.data.DataLoader(数据集, batch_size=128, shuffle=True)
print(testloader)

# 

经enumerate函数遍历，输出结果：

格式：[(索引, [tensor([128个图像数据]), tensor([128个对应标签])]),(索引, [tensor([...

# batch_idx :int 128                        | 当前运行的batch批次
# inputs    :torch.Size([128, 1, 28, 28])   | 图像
# targets   :torch.Size([128])              | 图像对应的类型标签
[(0, [tensor([[[[0., 0., 0.,  ..., 0., 0., 0.],
          [0., 0., 0.,  ..., 0., 0., 0.],
          [0., 0., 0.,  ..., 0., 0., 0.],
          ...,
          [0., 0., 0.,  ..., 0., 0., 0.],
          [0., 0., 0.,  ..., 0., 0., 0.],
          [0., 0., 0.,  ..., 0., 0., 0.]]],


        ...,
        [[[0., 0., 0.,  ..., 0., 0., 0.],
          [0., 0., 0.,  ..., 0., 0., 0.],
          [0., 0., 0.,  ..., 0., 0., 0.],
          ...,
          [0., 0., 0.,  ..., 0., 0., 0.],
          [0., 0., 0.,  ..., 0., 0., 0.],
          [0., 0., 0.,  ..., 0., 0., 0.]]]]), tensor([6, 8, 2, 0, 4, 3, 6, 3, 2, 5, 1, 3, 1, 3, 3, 2, 0, 6, 8, 2, 6, 7, 2, 5,
        2, 2, 8, 5, 7, 9, 3, 1, 0, 7, 5, 2, 7, 7, 8, 3, 7, 2, 2, 9, 8, 8, 7, 8,
        0, 0, 5, 3, 8, 1, 1, 0, 5, 4, 0, 4, 4, 5, 9, 1, 1, 7, 4, 9, 6, 9, 2, 0,
        7, 5, 8, 6, 4, 2, 0, 4, 2, 2, 6, 3, 1, 5, 7, 5, 4, 6, 1, 8, 8, 6, 4, 1,
        3, 3, 7, 2, 3, 1, 5, 9, 5, 1, 6, 6, 3, 6, 7, 4, 3, 4, 4, 1, 0, 3, 6, 4,
        6, 0, 3, 1, 4, 4, 7, 9])]),         # 以上为第0批次——128个数据的图像 标签
(1, [tensor([[[[0., 

• ---

  详细思路：

在我们训练深度学习网络时，经常需要使用不同的数据集训练网络。然而，网上的很多程序使用诸如mnist等通用数据集时，使用的多为专用的读取函数（例如：torchvision.datasets.MNIST()）。由于无法获知读取的数据结构，我们很难更换数据集。如果强行套用，会读取错位，导致无法输出正常的结果。

生成对抗网络 未正确设置输入结构，输出一堆马赛克

为此本人利用dataloader()函数、enumerate()函数，读取了其内部结构，并成功导入了自制的数据集。并将具体思路写成博客，希望对大家的研究有所帮助。

• 分析读取函数torchvision.datasets.MNIST()

transform = transforms.Compose([        # transforms.Compose：用于组合其他变换方式
    transforms.ToTensor(),              # transforms.ToTensor：将图片、数组转为tensor张量类型；转换从[0,255]->[0,1]
    transforms.Normalize([0.5], [0.5]), # transforms.Normalize：通过平均值和标准差来标准化一个tensor图像
])

trainset = torchvision.datasets.MNIST(root='./data', train=True, download=True, transform=transform)
# 下载mnist数据集(root：数据集目录、train：是否从training.pt创建数据集、download：从网上下载、transform：接收PIL图片并返回的转换函数)
trainloader = torch.utils.data.DataLoader(trainset, batch_size=128, shuffle=True)   
# DataLoader设置迭代对象，将Dataset分割为batch size大小的数据段、shuffle打乱数据

testset = torchvision.datasets.MNIST(root='./data', train=False, download=True, transform=transform)
# 下载mnist数据集(root：数据集目录、train：是否从test.pt创建数据集、download：从网上下载、transform：接收PIL图片并返回的转换函数)
testloader = torch.utils.data.DataLoader(testset, batch_size=100, shuffle=False)
'''
但是单独屏蔽两条中任意一条，下载的结果都一样。训练集、测试集是一起集成在4个文件里的？
'''

这段函数的作用，是在程序所在文件夹的./data文件夹分别读取训练集、测试集。如果没有便会从网上直接下载到文件夹。

下载的文件是4个压缩包，解压后是8个idx文件。

 

 然而这些idx文件无法用编辑器直接打开。

 本人曾想办法读取二进制文件，但并未成功。便直接在编辑器中显示其输入torch.utils.data.DataLoader()函数的值。

 读取作为输入值的testset

print(testset)

'''
Dataset MNIST
    Number of datapoints: 10000
    Root location: ./data
    Split: Test
    StandardTransform
Transform: Compose(
               ToTensor()
               Normalize(mean=[0.5], std=[0.5])
           )
'''


print(list(testset))     # 将数据以列表的形式显示


'''
[(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, 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]]]), 6),
          
(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, 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]]]), 7)]
'''

根据输出的中括号、小括号可知，数据集整体是一个列表[元组(tensor数据,int标签)]的格式。

• 读取数据、显示进度

确定了数据集的结构，就可以进行仿制了。基本思路就是将图片以tensor的格式与标签组成元组，而后将每张图片对应的元组组合成列表。

def read_file(data_path):
    fileNames = os.listdir(data_path)  # 获取当前路径下的文件名，返回List
    output = []
    for file in range(len(fileNames)):    # 读取文件夹内的所有图片及其名称
        img_folder = data_path + '\\' + fileNames[file]  #读取文件中的文件名
        image = cv2.imread(img_folder)
        im_gray = gray(image).view(-1, 36, 36)    # 将3通道图片转换为黑白单通道
        l = int(fileNames[file].split('_',1)[0])    # 图像名称格式:0_401  故以'_'分割图像名称，获取标签
        tup = (im_gray,l)   # 将数据与标签打包为元组
        output.append(tup)    # 将元组组合成列表

        if file%1000 == 999:        # 实时显示加载进度
            print("\rData loaded "+ str(file+1), end="  ")
    return output       #

• 实现mnist数据集生成

输出结果如下：

 由于，采用的是1000的数据集，训练生成的图像依然不够清晰，但经过20次迭代已依稀可以分辨数字的轮廓。

详细程序：

import re
import torch
import torch.nn as nn
import torch.nn.functional as F
import os
import cv2
from torchvision.utils import save_image
import matplotlib.pyplot as plt


# 损失函数
def loss_function(recon_x, x, mu, logvar):
    """
    :param recon_x: generated image
    :param x: original image
    :param mu: latent mean of z
    :param logvar: latent log variance of z
    """
    BCE_loss = nn.BCELoss(reduction='sum')
    reconstruction_loss = BCE_loss(recon_x, x)
    KL_divergence = -0.5 * torch.sum(1+logvar-torch.exp(logvar)-mu**2)
    #KLD_ele = mu.pow(2).add_(logvar.exp()).mul_(-1).add_(1).add_(logvar)
    #KLD = torch.sum(KLD_ele).mul_(-0.5)
    print(reconstruction_loss, KL_divergence)

    return reconstruction_loss + KL_divergence


# 构建VAE网络
class VAE(nn.Module):
    def __init__(self):
        super(VAE, self).__init__()
        self.fc1 = nn.Linear(1296, 400)      # 28*28,400
        self.fc2_mean = nn.Linear(400, 20)
        self.fc2_logvar = nn.Linear(400, 20)
        self.fc3 = nn.Linear(20, 400)
        self.fc4 = nn.Linear(400, 1296)

    def encode(self, x):
        h1 = F.relu(self.fc1(x))
        return self.fc2_mean(h1), self.fc2_logvar(h1)

    def reparametrization(self, mu, logvar):
        # sigma = 0.5*exp(log(sigma^2))= 0.5*exp(log(var))
        std = 0.5 * torch.exp(logvar)
        # N(mu, std^2) = N(0, 1) * std + mu
        z = torch.randn(std.size()) * std + mu
        return z

    def decode(self, z):
        h3 = F.relu(self.fc3(z))
        return torch.sigmoid(self.fc4(h3))

    def forward(self, x):
        mu, logvar = self.encode(x)
        z = self.reparametrization(mu, logvar)
        return self.decode(z), mu, logvar


# 转单通道
def gray(img):
    shape = img.shape
    height = shape[0]
    width = shape[1]
    dst = torch.empty(height,width)
    for h in range(0,height):
        for w in range(0,width):
            (b,g,r) = img[h,w]#*255
            gray = (int(r*313524) + int(g*615514) + int(b*119538))>>20     # 计算灰度
                # gray = (int(r) + int(g) + int(b)) // 3

            dst[h,w] = int(gray)
    
    return dst/255


# 读取图像、标签 并封装
def read_file(data_path):
    fileNames = os.listdir(data_path)  # 获取当前路径下的文件名，返回List
    output = []
    for file in range(len(fileNames)):    # 读取文件夹内的所有图片及其名称
        img_folder = data_path + '\\' + fileNames[file]  #读取文件中的文件名
        image = cv2.imread(img_folder)
        im_gray = gray(image).view(-1, 36, 36)    # 将3通道图片转换为黑白单通道
        l = int(fileNames[file].split('_',1)[0])    # 图像名称格式:0_401  故以'_'分割图像名称，获取标签
        tup = (im_gray,l)   # 将数据与标签打包为元组
        output.append(tup)    # 将元组组合成列表
    return output       #


batch_size = 128
train_num = 800
data_path = "图像地址"
mnist_dataset = read_file(data_path)


#加载MNIST数据集
trainset = mnist_dataset[:train_num]        # 切割数据集,划分 训练集、测试集
testset = mnist_dataset[train_num:]
trainloader = torch.utils.data.DataLoader(trainset, batch_size=batch_size, shuffle=True)
testloader = torch.utils.data.DataLoader(testset, batch_size=100, shuffle=False)


##### 训练网络
vae = VAE()
optimizer = torch.optim.Adam(vae.parameters(), lr=0.0003)
vae.train()
all_loss = 0.
# for batch_idx, (inputs, targets) in enumerate(trainloader):
#     print()
# Training
'''
enumerate() 函数用于将一个可遍历的数据对象(如列表、元组或字符串)组合为一个索引序列，同时列出数据和数据下标
'''
def train(epoch):
    vae.train()
    all_loss = 0.
    for batch_idx, (inputs, targets) in enumerate(trainloader):
        # print(batch_idx)
# batch_idx :int 128                        | 当前运行的batch批次
# inputs    :torch.Size([128, 1, 28, 28])   | 图像
# targets   :torch.Size([128])              | 图像对应的类型标签
        
        inputs, targets = inputs.to('cpu'), targets.to('cpu')
        # print(inputs.shape)     # torch.Size([128, 1, 28, 28])
        real_imgs = torch.flatten(inputs, start_dim=1)      # 将inputs里的第1维后的维度乘起来
        #print(real_imgs.shape)      # torch.Size([128, 784])

        # Train Discriminator
        gen_imgs, mu, logvar = vae(real_imgs)
        loss = loss_function(gen_imgs, real_imgs, mu, logvar)

        optimizer.zero_grad()
        loss.backward()
        optimizer.step()

        all_loss += loss.item()
    print('Epoch {}, loss: {:.6f}'.format(epoch, all_loss/(batch_idx+1)))
        # Save generated images for every epoch
    # print(gen_imgs.shape)       # torch.Size([96, 784])
    fake_images = gen_imgs.view(-1, 1,36, 36)     # view()函数，规整数据结构
    # print(fake_images.shape)        # torch.Size([96, 1, 28, 28]) 图片：8*12
    save_image(fake_images, 'MNIST_FAKE/fake_images-{}.png'.format(epoch + 1))
for epoch in range(20):
    train(epoch)


# 显示各迭代效果
data_path = "./MNIST_FAKE"      # 在本文件夹
fileNames = os.listdir(data_path)  # 获取当前路径下的文件名，返回List
fileNames.sort(key=lambda x:int(re.split('[.|_|-]',x)[2]))

print(fileNames)
for file in range(len(fileNames)):
    img_folder = data_path + '\\' + fileNames[file]  #文件中的文件名
    image = cv2.imread(img_folder)          
    num = (file)%5
    # print(num)
    if file%5 ==0:
        plt.figure(figsize=[20,4])
    
    plt.subplot(1,5,num+1)
    plt.imshow(image)
    plt.axis("off")
    plt.title(fileNames[file])
plt.show()



# 保存模型
torch.save(vae.state_dict(), './parameter/vae.pth')