【神经网络】Pytorch构建自己的训练数据集

​ 最近参加了一个比赛，需要对给定的图像数据进行分类，之前使用Pytorch进行神经网络模型的构建与训练过程中，都是使用的Pytorch内置的数据集，直接下载使用就好，而且下载下来的数据已经是Pytorch可以直接用于训练的Dataset类型，十分方便。但是如何使用Pytorch对自己的数据集进行训练，以及如何将数据集转换成Pytorch可以用于训练的Dataset数据类型，着实进行了一系列的摸索与尝试，特此进行记录。

​ 如果要重写Pytorch的Dataset类与DataLoader类，首先需要获取到自己的数据集，我的数据是存放在 “./data” 文件夹下面，其中 “./data/image/train/” 文件夹下面存放的是用于训练的图像数据，对应的图像标签存放在 “./data/train.labels.csv” 文件中，其中的数据存放格式是 “train_0.jpg 0” ，而 “./data/image/test/” 文件夹下面存放的是测试图像，不含标签，因此训练过程的验证集选择从训练图像中进行按照比例划分，选择使用sklearn中的对应函数，可以按照比例的进行随机划分。

​ 因此重写Pytorch的两个类，首先要将对应的图像与标签读取到，图像与标签进行匹配配对，然后进行后续的操作。

一、导入包：

import numpy as np
from sklearn.model_selection import train_test_split  # 将数据分为测试集和训练集
import os  # 打开读取文件
import cv2  # opencv读取图像
from matplotlib import pyplot as plt  # 测试图像是否读入进行绘制图像
import torch
from torch.autograd import Variable
from torchvision import transforms
from torch.utils.data import Dataset, DataLoader
from PIL import Image

二、读取图像文件夹：

def read_directory(directory_name, array_of_img, length, imgPath):
    # 参数为文件夹名称，图像数组，读取的文件长度，图像地址数组
    direct = os.listdir(r"./" + directory_name)
    # 乱序读取进来的数据集，对读取进来的数据集按照标号进行排序
    direct.sort(key=lambda x: int(x[6:-4]))
    # 如果输入-1，将其转换为长度加一，即将数据全部读入
    if length == -1:
        length = len(direct) + 1
    for filename in direct[0: length]:
        print("读取图像：" + filename)  # 测试是否按序读入
        iPath = "./" + directory_name + filename  # 读取当前图像的地址
        imgPath.append(iPath)  # 追加进图像地址数组中去
        img = cv2.imread(directory_name + "/" + filename)
        # 调整图像大小
        # img = cv2.resize(img, (32, 32))
        # 使用cv2.imread()接口读图像，读进来的是BGR格式以及【0～255】,所以需要将img转换为RGB格式正常显示
        img = img[:, :, [2, 1, 0]]
        # # 转为一维数组
        # img = np.array(img).flatten()
        # # print(img)
        array_of_img.append(img)

​ 注释已经基本上写的十分详细了，对代码进行简要的说明：

（1）读取到图像数据之后，将图像本身的数据与图像的地址数据分别存放在传入的列表"array_of_img"和"imgPath"中；

（2）读取图像数据之后，不知道Python读取进来的数据顺序是什么，但是有序的文件夹中的图像名称，读取进来之后就变成乱序了，因此使用正则表达式对图像的编号进行排序，保证其可以和对应的图像编号相对应；

（3）训练数据集中含有近八万张图像，如果每次测试都将图像全部读取进来，时间开销巨大，因此设置参数"length"，对读取文件夹中的图像列表进行切片，可以在训练中读取较少的图像用于测试。如果想要读取全部的参数，可以传入参数"-1"，但是Python的切片操作"[0:-1]“得到的列表不包含最后一个数据，因此对”-1"进行转换，设置为文件夹图像数据长度+1，就可以实现图像的全部读取进来；

（4）使用cv2.imread()接口读图像，读进来的是BGR格式以及[0～255]，如果对图像进行绘制的话，显示的颜色失真，所以需要将img转换为RGB格式正常显示。

三、显示图像：

def show_pic(images, length):
    i = 0
    plt.figure(figsize=(length, length))
    for img in images[0: length * length]:
        plt.subplot(length, length, i + 1)
        plt.imshow(img)
        # 关闭x，y轴显示
        plt.xticks([])
        plt.yticks([])
        i += 1

    plt.show()

​ 如果想要测试是否成功的读取到图像，可以选择性的调用这个函数对图像进行绘制，参数为图像本身的数据数组以及想要每行显示的图像数量（共绘制length*length个图像）。

四、读取图像对应标签：

def read_tag(tag_path, tags, length):
    # 参数为标签文件路径，标签数组，读取长度
    # 读取数据的标签
    with open(tag_path) as file:
        for line in list(file)[0:length]:
            # 读取前多少张图片的标签
            a = str(line).split("\t")
            b = a[1]  # 得到数字值
            tag = int(b[0])  # 去掉换行符
            tags.append(tag)

​ 同读取图像文件夹类似，需要对读取的标签数据进行格式检查，得到整型的图像标签，我的标签每行数据格式如"train_0.jpg 0"，因此这里的处理方法要根据自己的数据格式进行相应的调整。

五、分类图像数据集

#  对数据进行处理与加载
def data_load(img_path, tag_path, length, test_rate):
    # 获取数据与数据标签，并将数据分为训练集和测试集
    # 参数为图像文件夹，标签地址，length为读取的数据集长度，size为训练集测试集的区分比例

    images = []  # 图像数组
    tags = []  # 标签数组
    imPath = []  # 图像地址数组
    read_directory(img_path, images, length, imPath)
    read_tag(tag_path, tags, length)

    # 测试图像读取是否正常
    # show_pic(images, 10)

    # 转换为numpy数组
    images = np.array(images)
    tags = np.array(tags)

    # 返回处理得到的数据，有 X_train, X_test, y_train, y_test 四部分
    return train_test_split(images, tags, test_size=test_rate)

​ 如果使用的SkLearn等工具包对数据进行训练，可以直接对原始数据进行训练，则可以利用此函数读取相应的数据，并根据测试集的比例，随机将训练集划分为训练集与验证集。

六、生成文本文件：

def create_txt(img_path, tag_path, length, test_rate):
    # 生成训练集，验证集的数据地址与对应标签存储的数据对
    images = []  # 图像数组
    imagePath = []  # 图像地址数组
    tags = []  # 训练标签数组

    # 读取指定数量的图像，图像地址，以及对应的训练标签
    read_directory(img_path, images, length, imagePath)
    read_tag(tag_path, tags, length)

    # 根据比例随机划分训练集与测试集对应的图像地址和对应的标签
    X_train, X_test, y_train, y_test = train_test_split(imagePath, tags, test_size=test_rate)

    # w+意思是：
    # 打开一个文件用于读写。如果该文件已存在则打开文件，并从开头开始编辑，即原有内容会被删除。
    # 如果该文件不存在，创建新文件
    train = open('./data/train.txt', 'w+')
    test = open('./data/test.txt', 'w+')

    for i in range(len(X_train)):  # 生成训练文档
        name = X_train[i] + ' ' + str(int(y_train[i])) + '\n'
        train.write(name)
    for i in range(len(X_test)):  # 生成测试文档
        name = X_test[i] + ' ' + str(int(y_test[i])) + '\n'
        test.write(name)

    train.close()
    test.close()

​ 因为Pytorch中读取到的数据与标签类似于一个二元组，相互绑定，那么使用此函数将图像的地址与图像的标签进行绑定，生成一个txt文档，用于之后重写Dataset类，并且使用Sklearn的 “train_test_split()” 函数将数据根据比例随机划分为训练集与验证集，构建的txt文档格式为："./data/image/train/train_3.jpg 0"，用于之后的处理。

（注意使用的Python解释器，如果使用的是服务器上的解释器，生成的txt文档会在服务器的对应地址处，而本地没有这个txt文档，当初排查了好久。。。。。。）

七、重写Dataset类：

def default_loader(path):
    return Image.open(path).convert('RGB')


class MyDataset(Dataset):
    # 使用__init__()初始化一些需要传入的参数及数据集的调用
    def __init__(self, txt, resize, target_transform=None, loader=default_loader):
        super(MyDataset, self).__init__()
        # 对继承自父类的属性进行初始化
        fh = open(txt, 'r')
        imgs = []
        # 按照传入的路径和txt文本参数，以只读的方式打开这个文本
        for line in fh:  # 迭代该列表,按行循环txt文本中的内容
            line = line.strip('\n')
            line = line.rstrip('\n')
            # 删除本行string字符串末尾的指定字符，默认为空白符，包括空格、换行符、回车符、制表符
            words = line.split()
            # 用split将该行分割成列表，split的默认参数是空格，所以不传递任何参数时分割空格
            imgs.append((words[0], int(words[1])))
            # 把txt里的内容读入imgs列表保存
        self.imgs = imgs
        # 重新定义图像大小
        self.transform = transforms.Compose([transforms.Resize(size=(resize, resize)), transforms.ToTensor()])
        self.target_transform = target_transform
        self.loader = loader

    # 使用__getitem__()对数据进行预处理并返回想要的信息
    def __getitem__(self, index):
        fn, label = self.imgs[index]
        # fn是图片path
        img = self.loader(fn)
        # 按照路径读取图片
        if self.transform is not None:
            img = self.transform(img)
            # 数据标签转换为Tensor
        return img, label
        # return回哪些内容，那么在训练时循环读取每个batch时，就能获得哪些内容

    # 使用__len__()初始化一些需要传入的参数及数据集的调用
    def __len__(self):
        return len(self.imgs)

​ 除了注释中的说明外，因为不同的神经网络模型对于图像大小要求不一样，因此构建数据集时传入参数 “resize”，将图像进行大小调整为 “resize*resize” 大小的图像，用于模型训练。

八、创建DataLoader：

def MyDataLoader(img_path, tag_path, length, test_rate, Resize, batch_size):
    # 制作适用于torch的数据
    create_txt(img_path, tag_path, length, test_rate)
    train_data = MyDataset(txt='./data/train.txt', resize=Resize)
    test_data = MyDataset(txt='./data/test.txt', resize=Resize)
    train_loader = DataLoader(dataset=train_data, batch_size=batch_size, shuffle=True)
    test_loader = DataLoader(dataset=test_data, batch_size=batch_size, shuffle=False)
    return train_loader, test_loader

​ 首先构建图像与标签对应的文本数据，之后创建对应的dataset，然后构建成DataLoader用于Pytorch的模型训练。

九、测试：

def main():
    img_path = "data/image/train/"
    tag_path = "data/train.labels.csv"
    data_length = 10
    test_rate = 0.3
    batch_size = 2
    Resize = 64

    train_loader, test_loader = MyDataLoader(img_path, tag_path, data_length, test_rate, Resize, batch_size)
    for X, y in train_loader:
        print("Shape of X [N, C, H, W]: ", X.shape)
        print("Shape of y: ", y.shape, y.dtype)
        break


if __name__ == "__main__":
    main()

​ 输出为：

读取图像：train_0.jpg
读取图像：train_1.jpg
读取图像：train_2.jpg
读取图像：train_3.jpg
读取图像：train_4.jpg
读取图像：train_5.jpg
读取图像：train_6.jpg
读取图像：train_7.jpg
读取图像：train_8.jpg
读取图像：train_9.jpg
Shape of X [N, C, H, W]: torch.Size([2, 3, 64, 64])
Shape of y: torch.Size([2]) torch.int64

​ 整理不易，点赞支持~