Pytorch加载自己的数据集（以图片格式的Mnist数据集为例）
前言
初学pytorch，看了很多教程，发现所有教程在加载数据集的时候都用的pytorch已经定义好的模块，没有详细讲到如何使用Dataset和DataLoader加载自己格式多样的数据集，经过一段时间研究，成功跑通以图片为训练数据集的简单分类模型，现记录如下。
数据集在这里：
链接: https://pan.baidu.com/s/16T1IoAgOsepLqFRzjDck3g?pwd=h254 提取码: h254 复制这段内容后打开百度网盘手机App，操作更方便哦
一、数据集转换
Mnist是非常经典的数据集之一，从官网下载得到的是二进制的文件，与我们常用的图片格式不符，所以先将二进制文件转换为图像。
转换代码如下：
#-*- coding: utf-8-*-
import numpy as np
importstructimport os
import cv2
classDataUtils(object):
def__init__(self, filename=None, outpath=None):
self._filename= filename
self._outpath= outpath
self._tag='>' # 大端格式
self._twoBytes='II'
self._fourBytes='IIII'
self._pictureBytes='784B'
self._labelByte='1B'
self._twoBytes2= self._tag+ self._twoBytes
self._fourBytes2= self._tag+ self._fourBytes
self._pictureBytes2= self._tag+ self._pictureBytes
self._labelByte2= self._tag+ self._labelByte
self._imgNums=0
self._LabelNums=0
defgetImage(self):"""
将MNIST的二进制文件转换成像素特征数据"""
binfile=open(self._filename,'rb') # 以二进制方式打开文件
buf= binfile.read()
binfile.close()
index=0
numMagic, self._imgNums, numRows, numCols=struct.unpack_from(self._fourBytes2, buf, index)
index+=struct.calcsize(self._fourBytes)
images=[]print('image nums: %d'% self._imgNums)for i inrange(self._imgNums):
imgVal=struct.unpack_from(self._pictureBytes2, buf, index)
index+=struct.calcsize(self._pictureBytes2)
imgVal=list(imgVal)
images.append(imgVal)return np.array(images), self._imgNums
defgetLabel(self):"""
将MNIST中label二进制文件转换成对应的label数字特征"""
binFile=open(self._filename,'rb')
buf= binFile.read()
binFile.close()
index=0
magic, self._LabelNums=struct.unpack_from(self._twoBytes2, buf, index)
index+=struct.calcsize(self._twoBytes2)
labels=[]for x inrange(self._LabelNums):
im=struct.unpack_from(self._labelByte2, buf, index)
index+=struct.calcsize(self._labelByte2)
labels.append(im[0])return np.array(labels)
defoutImg(self, arrX, arrY, imgNums):"""
根据生成的特征和数字标号，输出图像"""
output_txt= self._outpath+'/img.txt'
output_file=open(output_txt,'a+')
m, n= np.shape(arrX)
# 每张图是28*28=784Bytefor i inrange(imgNums):
img= np.array(arrX[i])
img= img.reshape(28,28)#print(img)
outfile=str(i)+"_"+str(arrY[i])+".bmp"#print('saving file: %s'% outfile)
txt_line= outfile+" "+str(arrY[i])+'\n'
output_file.write(txt_line)
cv2.imwrite(self._outpath+'/'+ outfile, img)
output_file.close()if __name__=='__main__':
# 二进制文件路径，需要修改，和自己的相对应
trainfile_X='C:\\Users\\60058670\\Desktop\\MNIST\\train-images.idx3-ubyte'
trainfile_y='C:\\Users\\60058670\\Desktop\\MNIST\\train-labels.idx1-ubyte'
testfile_X='C:\\Users\\60058670\\Desktop\\MNIST\\t10k-images.idx3-ubyte'
testfile_y='C:\\Users\\60058670\\Desktop\\MNIST\\t10k-labels.idx1-ubyte'
# 加载mnist数据集
train_X, train_img_nums=DataUtils(filename=trainfile_X).getImage()
train_y=DataUtils(filename=trainfile_y).getLabel()
test_X, test_img_nums=DataUtils(testfile_X).getImage()
test_y=DataUtils(testfile_y).getLabel()
# 以下内容是将图像保存到本地文件中
path_trainset="C:\\Users\\60058670\\Desktop\\MNIST\\train"
path_testset="C:\\Users\\60058670\\Desktop\\MNIST\\test"if not os.path.exists(path_trainset):
os.mkdir(path_trainset)if not os.path.exists(path_testset):
os.mkdir(path_testset)DataUtils(outpath=path_trainset).outImg(train_X, train_y,int(train_img_nums/10)) #/10是只转换十分之一，用于测试DataUtils(outpath=path_testset).outImg(test_X, test_y,int(test_img_nums/10))二、构建自己的数据集
构建方法为继承Dataset类，用DataLoader加载
1.引入库
from torch.utils.data import Dataset
from torch.utils.data import DataLoader2.构建MnistDataset类
# 构建自己的数据集
classMnistDataset(Dataset):
def__init__(self, transform=None, lu_jing=None):
self.lu_jing= lu_jing
self.数据= os.listdir(self.lu_jing)
self.transform= transform
self.len=len(self.数据)
def__getitem__(self, index):
image_index= self.数据[index]
img_path= os.path.join(self.lu_jing, image_index)
img= Image.open(img_path)if self.transform:
img= self.transform(img)
label=int(image_index[-5])
label= self.oneHot(label)return img, label
def__len__(self):return self.len
# 将标签转为onehot编码
defoneHot(self, label):
tem= np.zeros(10)
tem[label]=1return torch.from_numpy(tem)3.搭建网络模型
只为演示，模型比较简单。
classModel(torch.nn.Module):
def__init__(self):super(Model, self).__init__()
self.Conv1= torch.nn.Conv2d(1,10, kernel_size=(5,5))
self.Conv2= torch.nn.Conv2d(10,20, kernel_size=(5,5))
self.pool= torch.nn.MaxPool2d(2)
self.fl= torch.nn.Linear(320,10)
defforward(self, x):
bs= x.size(0)
x= F.relu(self.pool(self.Conv1(x)))
x= F.relu(self.pool(self.Conv2(x)))
x= x.view(bs,-1)
x= self.fl(x)return x三 完整代码
import os
from PIL import Image
import torch
from torch.utils.data import Dataset
from torch.utils.data import DataLoader
import numpy as np
from torchvision import transforms
import torch.nn.functional as F
# 构建自己的数据集
classMnistDataset(Dataset):
def__init__(self, transform=None, lu_jing=None):
self.lu_jing= lu_jing
self.数据= os.listdir(self.lu_jing)
self.transform= transform
self.len=len(self.数据)
def__getitem__(self, index):
image_index= self.数据[index]
img_path= os.path.join(self.lu_jing, image_index)
img= Image.open(img_path)if self.transform:
img= self.transform(img)
label=int(image_index[-5])
label= self.oneHot(label)return img, label
def__len__(self):return self.len
# 将标签转为onehot编码
defoneHot(self, label):
tem= np.zeros(10)
tem[label]=1return torch.from_numpy(tem)
classModel(torch.nn.Module):
def__init__(self):super(Model, self).__init__()
self.Conv1= torch.nn.Conv2d(1,10, kernel_size=(5,5))
self.Conv2= torch.nn.Conv2d(10,20, kernel_size=(5,5))
self.pool= torch.nn.MaxPool2d(2)
self.fl= torch.nn.Linear(320,10)
defforward(self, x):
bs= x.size(0)
x= F.relu(self.pool(self.Conv1(x)))
x= F.relu(self.pool(self.Conv2(x)))
x= x.view(bs,-1)
x= self.fl(x)return xif __name__=='__main__':
# 训练集路径
train_data="C:\\Users\\60058670\\Desktop\\MNIST\\train"
transform= transforms.Compose([transforms.ToTensor()]) # 归一化处理
data=MnistDataset(transform=transform, lu_jing=train_data)
data_loader=DataLoader(data, batch_size=200, shuffle=True) # 使用DataLoader加载数据
model=Model()
criterion= torch.nn.CrossEntropyLoss() # 交叉熵损失
optimizer= torch.optim.SGD(model.parameters(), lr=0.01, momentum=0.5) # model.parameters()自动完成参数的初始化操作for epoch inrange(20):for i, data1 inenumerate(data_loader,0): # train_loader 是先shuffle后mini_batch
inputs, labels= data1
y_pred=model(inputs)
loss=criterion(y_pred, labels)
optimizer.zero_grad()
loss.backward()
optimizer.step()if epoch%5==0:print(epoch, loss.item())
# 测试集路径
test_data='C:\\Users\\60058670\\Desktop\\MNIST\\test'
x_test=MnistDataset(transform=transform, lu_jing=test_data)
x_test=DataLoader(x_test, batch_size=100, shuffle=False) # 使用DataLoader加载数据
total=0
correct=0for i, data inenumerate(x_test,0): # train_loader 是先shuffle后mini_batch
inputs, labels= data
y_pred=model(inputs)
_, labels= torch.max(labels.data, dim=1)
_, predicted= torch.max(y_pred.data, dim=1)
total+= labels.size(0)
correct+=(predicted== labels).sum().item()print('accuracy on test set: {} % '.format(100* correct/ total))print(correct, total)总结
纸上得来终觉浅，绝知此事要躬行。自己动手写了代码就会发现一堆问题，知识就是在解决问题的过程中积累的。初学不久，有问题大家可以一起交流讨论。