pytorch实现建立自己的数据集(以mnist为例)

本文将原始的numpy array数据在pytorch下封装为Dataset类的数据集，为后续深度网络训练提供数据。

加载并保存图像信息

首先导入需要的库，定义各种路径。

import osimport matplotlibfrom keras.datasets import mnistimport numpy as npfrom torch.utils.data.dataset import Datasetfrom PIL import Imageimport scipy.miscroot_path = 'E:/coding_ex/pytorch/Alexnet/data/'base_path = 'baseset/'training_path = 'trainingset/'test_path = 'testset/'

这里将数据集分为三类，baseset为所有数据（trainingset+testset），trainingset是训练集，testset是测试集。直接通过keras.dataset加载mnist数据集，不能自动下载的话可以手动下载.npz并保存至相应目录下。

def LoadData(root_path, base_path, training_path, test_path):  (x_train, y_train), (x_test, y_test) = mnist.load_data()  x_baseset = np.concatenate((x_train, x_test))  y_baseset = np.concatenate((y_train, y_test))  train_num = len(x_train)  test_num = len(x_test)    #baseset  file_img = open((os.path.join(root_path, base_path)+'baseset_img.txt'),'w')  file_label = open((os.path.join(root_path, base_path)+'baseset_label.txt'),'w')  for i in range(train_num + test_num):    file_img.write(root_path + base_path + 'img/' + str(i) + '.png/n') #name    file_label.write(str(y_baseset[i])+'/n') #label#    scipy.misc.imsave(root_path + base_path + '/img/'+str(i) + '.png', x_baseset[i])    matplotlib.image.imsave(root_path + base_path + 'img/'+str(i) + '.png', x_baseset[i])  file_img.close()  file_label.close()    #trainingset  file_img = open((os.path.join(root_path, training_path)+'trainingset_img.txt'),'w')  file_label = open((os.path.join(root_path, training_path)+'trainingset_label.txt'),'w')  for i in range(train_num):    file_img.write(root_path + training_path + 'img/' + str(i) + '.png/n') #name    file_label.write(str(y_train[i])+'/n') #label#    scipy.misc.imsave(root_path + training_path + '/img/'+str(i) + '.png', x_train[i])    matplotlib.image.imsave(root_path + training_path + 'img/'+str(i) + '.png', x_train[i])  file_img.close()  file_label.close()    #testset  file_img = open((os.path.join(root_path, test_path)+'testset_img.txt'),'w')  file_label = open((os.path.join(root_path, test_path)+'testset_label.txt'),'w')  for i in range(test_num):    file_img.write(root_path + test_path + 'img/' + str(i) + '.png/n') #name    file_label.write(str(y_test[i])+'/n') #label#    scipy.misc.imsave(root_path + test_path + '/img/'+str(i) + '.png', x_test[i])    matplotlib.image.imsave(root_path + test_path + 'img/'+str(i) + '.png', x_test[i])  file_img.close()  file_label.close()

使用这段代码时，需要建立相应的文件夹及.txt文件，./data文件夹结构如下：

/img文件夹

由于mnist数据集其实是灰度图，这里用matplotlib保存的图像是伪彩色图像。