超简单！用pytorch训练自己的数据集

在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')

全预测对啦

有问题可以留言相互交流哦。