在pytorch官网上的tutoral中，教程给了一个例子，训练一个分类器，下载CIFAR10数据集，用一个简单的CNNC网络训练模型。在我们下载的CIFAR10数据集中，已标签的数据被打包并封装在data文件夹里。
我们可以在这个例子的基础上，把自己的数据放到其模型下，并实现一个简易的分类器。
首先是配置pytorch的各种环境，有很多教程的，他们讲的又详细又好，我就不多讲啦。
先导入各种模块
#encoding=utf-8 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 np
如果报错就是各种库各种包没有配好。没有报错就进行下一步啦
首先定义我们的训练数据。我们采用的是ImageFolder函数。将我们要训练的数据放在一个路径为path的文件夹下，假如这个文件夹名为train。在我们的train文件夹下有N个子文件夹，每个子文件夹代表一个分类，一共有N类。
transforms.Compose()是把几个transform语句合并到一起
transforms.Resize()可以传入两类参数。一类是（h,w）这将把图片缩放到（h,w）大小（长宽比会改变）。另一类是传入单个参数，这将把图片经过缩放后（保持长宽比不变），将最短的边缩放到传入的参数。
transforms.ToTensor()是数据类型的转换。我们得到的trainset就和官网tutorals里的trainset有一样的格式啦。
需要注意的是我们自己的trainset的图像大小必须为32×32，因为全连接层是固定尺寸的输入输出，所以在卷基层之前的输入，也就是我们训练的图片尺寸大小要求是固定的。CIFAR10的大小是32×32，所以自己的数据也要被缩放到32×32。如果需要其他尺寸的，可以修改CNN的内部结构。
def loadtraindata(): path = r"/home/********/folder/train" # 路径 trainset = torchvision.datasets.ImageFolder(path, transform=transforms.Compose([ transforms.Resize((32, 32)), # 将图片缩放到指定大小（h,w）或者保持长宽比并缩放最短的边到int大小 transforms.CenterCrop(32), transforms.ToTensor()]) ) trainloader = torch.utils.data.DataLoader(trainset, batch_size=4, shuffle=True, num_workers=2) return trainloader
接着定义CNN网络。我这里用的结构完全照搬tutorals里的结构。注意我们这里同样是10个输出，也就是10种分类。如果你只有9种分类记得改变参数哦。
class Net(nn.Module): # 定义网络，继承torch.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))) x = x.view(-1, 16 * 5 * 5) # .view( )是一个tensor的方法，使得tensor改变size但是元素的总数是不变的。 # 从卷基层到全连接层的维度转换 x = F.relu(self.fc1(x)) x = F.relu(self.fc2(x)) x = self.fc3(x) return
我用的数据库是车牌里的数字，从0～9一共10类，所以定义的种类还是10种。
classes = ('0','1', '2', '3', '4', '5', '6', '7', '8', '9')
测试集也和训练集相似。我这里将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)), # 将图片缩放到指定大小（h,w）或者保持长宽比并缩放最短的边到int大小 transforms.ToTensor()]) ) testloader = torch.utils.data.DataLoader(testset, batch_size=25, shuffle=True, num_workers=2) return testloader
接着我们需要训练这个CNN网络。用trainandsave函数来训练CNN的参数并用torch.save保存其参数。参数文件.pkl会放在py文件的同一个路径下。
def trainandsave(): trainloader = loadtraindata() # 神经网络结构 net = Net() optimizer = optim.SGD(net.parameters(), lr=0.001, momentum=0.9) # 学习率为0.001 criterion = nn.CrossEntropyLoss() # 损失函数也可以自己定义，我们这里用的交叉熵损失函数 # 训练部分 for epoch in range(5): # 训练的数据量为5个epoch，每个epoch为一个循环 # 每个epoch要训练所有的图片，每训练完成200张便打印一下训练的效果（loss值） running_loss = 0.0 # 定义一个变量方便我们对loss进行输出 for i, data in enumerate(trainloader, 0): # 这里我们遇到了第一步中出现的trailoader，代码传入数据 # enumerate是python的内置函数，既获得索引也获得数据 # get the inputs inputs, labels = data # data是从enumerate返回的data，包含数据和标签信息，分别赋值给inputs和labels # wrap them in Variable inputs, labels = Variable(inputs), Variable(labels) # 转换数据格式用Variable optimizer.zero_grad() # 梯度置零，因为反向传播过程中梯度会累加上一次循环的梯度 # forward + backward + optimize outputs = net(inputs) # 把数据输进CNN网络net loss = criterion(outputs, labels) # 计算损失值 loss.backward() # loss反向传播 optimizer.step() # 反向传播后参数更新 running_loss += loss.data[0] # loss累加 if i % 200 == 199: print('[%d, %5d] loss: %.3f' % (epoch + 1, i + 1, running_loss / 200)) # 然后再除以200，就得到这两百次的平均损失值 running_loss = 0.0 # 这一个200次结束后，就把running_loss归零，下一个200次继续使用 print('Finished Training') # 保存神经网络 torch.save(net, 'net.pkl') # 保存整个神经网络的结构和模型参数 torch.save(net.state_dict(), 'net_params.pkl') # 只保存神经网络的模型参数
如果我们只进行训练阶段，只需要执行
trainandsave()
就可以了。执行完毕后会保存两个参数文件。如果已经训练好了，就不需要执行训练用的这个函数了。接下来是预测。先定义两个函数：
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() # imshow(torchvision.utils.make_grid(images,nrow=5)) # nrow是每行显示的图片数量，缺省值为8 print('GroundTruth: ' , " ".join('%5s' % classes[labels[j]] for j in range(25))) # 打印前25个GT（test集里图片的标签） outputs = net(Variable(images)) _, predicted = torch.max(outputs.data, 1) print('Predicted: ', " ".join('%5s' % classes[predicted[j]] for j in range(25))) # 打印前25个预测值
如果我们只进行训练阶段，首先需要一个根目录里已经训练好了的net.pkl文件，再执行
test()
最后显示
('GroundTruth: ', ' 5 0 7 1 2 1 5 9 8 2 6 3 2 0 4 9 3 9 8 7 1 7 6 4 4')
('Predicted: ', ' 5 0 7 1 2 1 5 9 8 2 6 3 2 0 4 9 3 9 8 7 1 7 6 4 4')
全预测对啦
有问题可以留言相互交流哦。