Pytorch：数据集构建、加载、划分

在利用神经网络对自己的数据进行分析时，首先要对数据进行处理，构建用来训练和测试的数据集，并对其进行加载。具体方法如下：

本文中，我们将几类类别不同的数据分别放在不同的文件夹中，并将所有类别的文件夹放在一个大文件夹里，为构建数据集，我们首先根据文件路径和对应文件的标签生成一个.txt文件：

import os
def generate(dir, label):
    files = os.listdir(dir)
    files.sort()
    listText = open('all_data.txt', 'a')
    for file in files:
        fileType = os.path.split(file)
        if fileType[1] == '.txt':
            continue
        file1 = os.path.join(dir, file)
        print(file1)
        name = file1 + ' ' + str(int(label)) + '\n'
        listText.write(name)
    listText.close()


outer_path = 'D:/'  #存储数据的大文件夹路径

if __name__ == '__main__':
    i = 0
    folderlist = os.listdir(outer_path)  # 列举所有类别文件夹
    for folder in folderlist:
        generate(os.path.join(outer_path, folder), i)