机器学习11：pytorch训练自定义数据集简单示例
本文整理总结自博客用portorch训练自己的数据集，在pytorch官网例程的基础上将自己的数据放到其模型下，实现一个识别手写数字的简易分类器。
1.环境配置及模块导入：
首先配置pytorch的运行环境，然后导入各种模块。
import torch
import torchvision
import torchvision.transforms as transforms
import torch.nn as nn
import torch.nn.functional as F
from torch.autograd import Variable
import torch.optim as optim
import matplotlib.pyplot as plt
import numpy as np2.处理自定义数据集
（1）处理训练集
首先调用ImageFolder()函数处理训练数据。
将我们要训练的数据放在一个路径为path的文件夹下，假如这个文件夹名为train。在我们的train文件夹下有N个子文件夹，每个子文件夹代表一个分类，一共有N类。
transforms.Compose()的功能是合并几个transform语句。
transforms.Resize()的功能是缩放图片的尺寸。Resize()函数有两种参数设置方式，第一种是Resize（h,w），将图片缩放到（h,w）大小，这种方式会改变图像的长宽比；第二种是Resize（h），这种方式不改变图像的长宽比，使图片经过缩放后最短的边长度等于传入的参数h。
transforms.CenterCrop(size)的功能是依据给定的size从中心裁剪。
transforms.ToTensor()的功能是转换数据类型，使得到的自定义数据集trainset和官网例程中的trainset格式相同。
def loadtraindata():
# 设置路径
path = r"/home/********/folder/train"
# 将图片缩放到指定大小
trainset = torchvision.datasets.ImageFolder(path,transform=transforms.Compose([
transforms.Resize((32, 32)),
transforms.CenterCrop(32),
transforms.ToTensor()]))
trainloader = torch.utils.data.DataLoader(trainset, batch_size=4,shuffle=True,
num_workers=2)
return trainloader本文用到的模型结构完全照搬官网例程，因为全连接层是固定尺寸的输入输出，所以trainset的图像大小也设置为与例程相同的32×32。如果需要其他尺寸的输入需要修改网络模型的内部结构。
（2）处理测试集
测试集和训练集相似，这里将batch_size改成了25，在测试阶段每次预测25张图片并以5×5的排列显示。
def loadtestdata():
path = r"/home/********/folder/test"
testset = torchvision.datasets.ImageFolder(path,transform=transforms.Compose([
transforms.Resize((32, 32)),
transforms.ToTensor()])
)
testloader = torch.utils.data.DataLoader(testset, batch_size=25,
shuffle=True, num_workers=2)
return testloader
3.定义网络结构
本文用到的模型结构完全照搬官网例程，注意在这里输出是10种分类，与例程保持一致。如果分类数不同需要修改网络结构的参数。
class Net(nn.Module):
def __init__(self):
super(Net, self).__init__()
self.conv1 = nn.Conv2d(3, 6, 5) # 卷积层
self.pool = nn.MaxPool2d(2, 2) # 池化层
self.conv2 = nn.Conv2d(6, 16, 5) # 卷积层
self.fc1 = nn.Linear(16 * 5 * 5, 120) # 全连接层
self.fc2 = nn.Linear(120, 84)
self.fc3 = nn.Linear(84, 10) # 10个输出
def forward(self, x): # 前向传播
x = self.pool(F.relu(self.conv1(x))) # F就是torch.nn.functional
x = self.pool(F.relu(self.conv2(x)))
# .view( )是一个tensor方法，使得tensor改变size但是元素的总数保持，实现从卷基层到全连接层的维度转换
x = x.view(-1, 16 * 5 * 5)
x = F.relu(self.fc1(x))
x = F.relu(self.fc2(x))
x = self.fc3(x)
return 本文数字识别的分类情况为classes = ('0','1', '2', '3', '4','5', '6', '7', '8', '9')，分类数为10类，与例程相同。
4.训练网络
用trainandsave()函数来训练CNN的参数并用torch.save保存其参数，保存后参数文件.pkl会放在py文件的同一个路径下。
如果只进行训练阶段，那么只需要执行trainandsave()函数。
def trainandsave():
trainloader = loadtraindata()
# 【1】定义神经网络结构
net = Net()
optimizer = optim.SGD(net.parameters(), lr=0.001, momentum=0.9) # 学习率为0.001
criterion = nn.CrossEntropyLoss() # 损失函数（交叉熵）
# 【2】训练部分
# 训练的数据量为5个epoch，每个epoch为一个循环
for epoch in range(5):
# 每个epoch训练所有图片，每训练完成200张打印一次训练损失值查看训练效果
# 定义一个变量方便对loss进行输出
running_loss = 0.0
# 这里用到trailoader，代码传入数据
# enumerate是python的内置函数，既获得索引也获得数据
for i, data in enumerate(trainloader, 0):
# data是从enumerate返回的data，包含数据和标签信息，分别赋值给inputs和labels
inputs, labels = data
# 转换数据格式用Variable
inputs, labels = Variable(inputs), Variable(labels)
# 梯度置零，因为反向传播过程中梯度会累加上一次循环的梯度
optimizer.zero_grad()
# 把数据输进网络net
outputs = net(inputs)
# 计算损失值
loss = criterion(outputs, labels)
# loss反向传播
loss.backward()
# 反向传播后参数更新
optimizer.step()
# 损失累加
running_loss += loss.data[0]
if i % 200 == 199:
# 除以200，获得两百次的平均损失值
print('[%d, %5d] loss: %.3f' %(epoch + 1, i + 1, running_loss / 200)) # 这一个200次结束后，就把running_loss归零，下一个200次继续使用
running_loss = 0.0
print('Finished Training')
# 【3】保存神经网络
torch.save(net, 'net.pkl') # 保存整个神经网络的结构和模型参数
torch.save(net.state_dict(), 'net_params.pkl') # 只保存神经网络的模型参数
5.测试网络
执行测试阶段首先需要一个根目录里已经训练好了的net.pkl文件，再执行test()函数。
def reload_net():
trainednet = torch.load('net.pkl')
return trainednet
def imshow(img):
img = img / 2 + 0.5 # unnormalize
npimg = img.numpy()
plt.imshow(np.transpose(npimg, (1, 2, 0)))
plt.show()在imshow()中的transpose()的作用可以简单理解为格式转换，不进行格式转换的数据无法正常读取。
def test():
testloader = loadtestdata()
net = reload_net()
dataiter = iter(testloader)
images, labels = dataiter.next()
# nrow为每行显示的图片数量
imshow(torchvision.utils.make_grid(images,nrow=5))
# 打印前25个test集里图片的标签
print('GroundTruth: ', " ".join('%5s' % classes[labels[j]] for j in range(25)))
outputs = net(Variable(images))
_, predicted = torch.max(outputs.data, 1)
# 打印前25个预测值
print('Predicted: ', " ".join('%5s' % classes[predicted[j]] for j in range(25)))
