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

时间:2024-03-31 13:43:17

在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,代码传入数据
            # enumeratepython的内置函数,既获得索引也获得数据
            # get the inputs
            inputs, labels = data  # data是从enumerate返回的data,包含数据和标签信息,分别赋值给inputslabels

            # 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()

最后显示

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

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

全预测对啦

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