我的实践:通过蚂蚁、蜜蜂二分类问题了解如何基于Pytorch构建分类模型

1.数据集准备

本例采用了pytorch教程提供的蜜蜂、蚂蚁二分类数据集(点击可直接下载)。该数据集的文件夹结构如下图所示。这里面有些黑白的照片,我把它们删掉了,因为黑白照片的通道数是1,会造成Tensor的维度不一致。可以看出数据集分为训练集和测试集,训练集用于训练模型,测试集用于测试模型的泛化能力。在训练集和测试集下又包含了"ants"和"bees"两个文件夹,这两个文件夹的名称即图片的标签,在加载数据的时候需要用到这一点。有了数据,我们就想办法把这些数据处理成pytorch框架下的Dataset需要的格式。

请添加图片描述

2.pytorch Dataset 处理图片数据

pytorch为我们处理数据提供了一个模板,这个模板就是Dataset,我们在处理数据时继承这个类。在处理数据时要注意以下几点:

  1. 可以用PIL的Image加载图片,但要将图片处理成tensor,而且tensor的维度要一致。这是因为nn模型的输入都是tensor格式,而且要求一个batchsize的tensor维度是一样的。实现上述可能可以使用torchvision的transforms。由于我用的CPU训练模型,所以对图片压缩的比较厉害,全压缩成33232的图片了。
  2. "ants"和"bees"两个文件夹的名称就是图片的标签,但是getitem的返回值应该是一个值。在这里"ants"标签返回0,"bees"标签返回1。
  3. 看数据的预处理对不对,可以用一段代码测试一下,将数据加载到DataLoader,然后循环取出数据,并把这些数据及其标签打印出来,或者记录到tensorboard上去,看每一次迭代返回的数据是否和自己预想的一样。

下面是代码,保存在dataProcess.py文件中。

rom torch.utils.data import Dataset
from torch.utils.data import DataLoader
from PIL import Image
import os
from torchvision import transforms
from torch.utils.tensorboard import SummaryWriter

class MyData(Dataset):
    # 把图片所在的文件夹路径分成两个部分,一部分是根目录,一部分是标签目录,这是因为标签目录的名称我们需要用到
    def __init__(self, root_dir, label_dir):
        self.root_dir = root_dir
        self.label_dir = label_dir
        # 图片所在的文件夹路径由根目录和标签目录组成
        self.path = os.path.join(self.root_dir, self.label_dir)
        # 获取文件夹下所有图片的名称
        self.img_names = os.listdir(self.path)

    def __getitem__(self, idx):
        img_name = self.img_names[idx]
        img_item_path = os.path.join(self.root_dir, self.label_dir, img_name)
        img = Image.open(img_item_path)
        # 将图片处理成Tensor格式,并将维度设置成32*32的
        # 图片的维度可能不一致,这里一定要用resize统一一下,否则会出错
        trans = transforms.Compose(
            [
                transforms.ToTensor(),
                transforms.Resize((32, 32))
            ])
        img_tensor = trans(img)
        # 根据标签目录的名称来确定图片是哪一类,如果是"ants",标签设置为0,如果是"bees",标签设置为1
        # 这个地方要注意,我们在计算loss的时候用交叉熵nn.CrossEntropyLoss()
        # 交叉熵的输入有两个,一个是模型的输出outputs,一个是标签targets,注意targets是一维tensor
        # 例如batchsize如果是2,ants的targets的应该[0,0],而不是[[0][0]]
        # 因此label要返回0,而不是[0]
        label = 0 if self.label_dir == "ants" else 1
        return img_tensor,  label

    def __len__(self):
        return len(self.img_names)

# 用下面这段代码测试一下加载数据有没有问题
if __name__ == "__main__":
    # 注意hymenoptera_data和代码在同一级目录
    root_dir = "hymenoptera_data/train"
    ants_label = "ants"
    bees_label = "bees"
    # 蚂蚁数据集
    ants_dataset = MyData(root_dir, ants_label)
    # 蜜蜂数据集
    bees_dataset = MyData(root_dir, bees_label)
    # 蚂蚁数据集和蜜蜂数据集合并
    train_dataset = ants_dataset + bees_dataset
    # 利用dataLoader加载数据集
    train_dataloader = DataLoader(dataset=train_dataset, batch_size=64, shuffle=True)
    # tensorboard的writer
    writer = SummaryWriter("logs")
    for step, train_data in enumerate(train_dataloader):
        imgs, targets = train_data
        # 每迭代一次就把一个batch的图片记录到tensorboard
        writer.add_images("test", imgs, step)
        # 每迭代一次就把一个batch的图片标签打印出来
        print(targets)
    writer.close()

在测试时tensorboard记录的信息在logs文件夹,在terminal输入tensorboard --logdir=logs启动tensorboard,将tensorboard给出的网址输入到网页,可以看到每一个batch的图片。下图展示了第一个batch的图片。可以看到,取出了64张图片,和batchsize=64是对应的。另外可以看到,把图片压缩成32*32后,确实很模糊了,人眼都很难看出哪个是蚂蚁,哪个是蜜蜂。


请添加图片描述

下面这个图展示了第一个batch所有图片的标签,0表示蚂蚁,1表示蜜蜂,仔细看一下图片和标签应该是对应的。


请添加图片描述

3.网络模型设计

我们把图片处理成3*32*32的tensor了,用如下图所示的卷积神经网络模型。第一层卷积网络采用5*5的卷积核,stride=1,pading=2。第一层卷积的代码是:nn.Conv2d(3, 32, 5, 1, 2),第一个参数3是输入的通道数,第二个参数32是输出的通道数,第三个参数5是卷积核的大小,第四个参数1是stride,第五个参数2是padding。


卷积网络模型.png

输出高H,和宽度W计算公式如下所示(注意dilation默认为0)。



因此,通过第一层卷积后,高度H为,

同理宽度W也为32。所以输出的大小就32*32*32。接下来,再用一个max-Pooling进行一次池化,池化核的大小是2*2。该池化层的代码是nn.MaxPool2d(2)。池化输出高H,和宽度W计算公式和卷积计算方式一摸一样。在默认的情况下,stride和池化和的大小一样,pading=0,dilation=0。所以第一次池化后,输出的高度H为,

同理,输出的宽度H为16。因此,输出的维度是32*16*16。
后面的输出维度计算方式同上,不再罗嗦了。然后再通过两次卷积和两次池化,后面的输出维度计算方式同上,不再罗嗦了,最终得到一个维度为64*4*4的特征。在做分类之前,首先要把这个三维Tensor拉直成一维Tensor,代码是nn.Flatten()。拉直之后的一维Tensor大小就是。最后通过一个全连接层完成分类任务,全连接层的输入大小是1024,输出的大小是类别的个数,即2,代码是nn.Linear(64 * 4 * 4, 2)。

当完成所有模型的构建后,可以用一段代码来测试一下模型是否有误。例如这里模型的输入在[3,32,32]Tensor的基础上,还需要再增加一维batchsize,所以输入的维度应该是[batchsize,3,32,32]。我们可以生成一个这样维度的数据,例如假设batchsize=3,可以这样生成一个输入:x = torch.ones((3, 3, 32, 32))。然后把x送给模型,看模型是否能正常输出,输出的维度是否是我们预期的。我们还可以借助于Tensorboard来将模型可视化,通过界面把模型展开,看是否正确。
下面是所有的代码,保存在model.py文件中。

from torch import nn
import torch
from torch.utils.tensorboard import SummaryWriter

class MyModel(nn.Module):
    def __init__(self):
        super(MyModel, self).__init__()
        self.model = nn.Sequential(
            nn.Conv2d(3, 32, 5, 1, 2),
            nn.MaxPool2d(2),
            nn.Conv2d(32, 32, 5, 1, 2),
            nn.MaxPool2d(2),
            nn.Conv2d(32, 64, 5, 1, 2),
            nn.MaxPool2d(2),
            nn.Flatten(),
            nn.Linear(64 * 4 * 4, 2)
        )

    def forward(self, x):
        x = self.model(x)
        return x

# 这段代码测试model是否正确
if __name__ == "__main__":
    my_model = MyModel()
    x = torch.ones((3, 3, 32, 32))
    y = my_model(x)
    print(y.shape)
    # 利用tensorboard可视化模型
    writer = SummaryWriter("graph_logs")
    writer.add_graph(my_model, x)
    writer.close()

模型测试代码打印的输出维度是[3,2],3是batchsize,2是全连接层最后的输出维度,和类别的个数是一致的。利用Tensorboard将模型可视化后,如下图所示,还可以进一步展开。


请添加图片描述

4.模型的训练与测试

模型的训练与测试就不细讲了,和其他模型训练的套路一样的,基本思路可以看我的第一篇[pytorch入门文章](我的实践:通过一个简单线性回归入门pytorch - (jianshu.com)
)。下面直接给出代码。

from model import *
from dataProcess import *
import matplotlib.pyplot as plt
import time

# 加载训练数据
train_root_dir = "hymenoptera_data/train"
train_ants_label = "ants"
train_bees_label = "bees"
train_ants_dataset = MyData(train_root_dir, train_ants_label)
train_bees_dataset = MyData(train_root_dir, train_bees_label)
train_dataset = train_ants_dataset + train_bees_dataset
train_data_loader = DataLoader(dataset=train_dataset, batch_size=128, shuffle=True)
train_data_len = len(train_dataset)
# 加载测试数据
test_root_dir = "hymenoptera_data/val"
test_ants_label = "ants"
test_bees_label = "bees"
test_ants_dataset = MyData(test_root_dir, test_ants_label)
test_bees_dataset = MyData(test_root_dir, test_bees_label)
test_dataset = test_ants_dataset + test_bees_dataset
test_data_loader = DataLoader(dataset=test_dataset, batch_size=256, shuffle=True)
test_data_len = len(test_dataset)
print(f"训练集长度:{train_data_len}")
print(f"测试集长度:{test_data_len}")
# 创建网络模型
my_model = MyModel()

# 损失函数
loss_fn = nn.CrossEntropyLoss()

# 优化器
learning_rate = 5e-3
optimizer = torch.optim.SGD(my_model.parameters(), lr=learning_rate)
# Adam 参数betas=(0.9, 0.99)
# optimizer = torch.optim.Adam(my_model.parameters(), lr=learning_rate, betas=(0.9, 0.99))
# 总共的训练步数
total_train_step = 0
# 总共的测试步数
total_test_step = 0
step = 0
epoch = 500

writer = SummaryWriter("logs")
train_loss_his = []
train_totalaccuracy_his = []
test_totalloss_his = []
test_totalaccuracy_his = []
start_time = time.time()
my_model.train()
for i in range(epoch):
    print(f"-------第{i}轮训练开始-------")
    train_total_accuracy = 0
    for data in train_data_loader:
        imgs, targets = data
        writer.add_images("tarin_data", imgs, total_train_step)
        output = my_model(imgs)
        loss = loss_fn(output, targets)
        train_accuracy = (output.argmax(1) == targets).sum()
        train_total_accuracy = train_total_accuracy + train_accuracy
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()
        total_train_step = total_train_step + 1
        train_loss_his.append(loss)
        writer.add_scalar("train_loss", loss.item(), total_train_step)
    train_total_accuracy = train_total_accuracy / train_data_len
    print(f"训练集上的准确率:{train_total_accuracy}")
    train_totalaccuracy_his.append(train_total_accuracy)
    # 测试开始
    total_test_loss = 0
    my_model.eval()
    test_total_accuracy = 0
    with torch.no_grad():
        for data in test_data_loader:
            imgs, targets = data
            output = my_model(imgs)
            loss = loss_fn(output, targets)
            total_test_loss = total_test_loss + loss
            test_accuracy = (output.argmax(1) == targets).sum()
            test_total_accuracy = test_total_accuracy + test_accuracy
        test_total_accuracy = test_total_accuracy / test_data_len
        print(f"测试集上的准确率:{test_total_accuracy}")
        print(f"测试集上的loss:{total_test_loss}")
        test_totalloss_his.append(total_test_loss)
        test_totalaccuracy_his.append(test_total_accuracy)
        writer.add_scalar("test_loss", total_test_loss.item(), i)
end_time = time.time()
total_train_time = end_time-start_time
print(f'训练时间: {total_train_time}秒')
writer.close()
plt.plot(train_loss_his, label='Train Loss')
plt.legend(loc='best')
plt.xlabel('Steps')
plt.show()
plt.plot(test_totalloss_his, label='Test Loss')
plt.legend(loc='best')
plt.xlabel('Steps')
plt.show()

plt.plot(train_totalaccuracy_his, label='Train accuracy')
plt.plot(test_totalaccuracy_his, label='Test accuracy')
plt.legend(loc='best')
plt.xlabel('Steps')
plt.show()

通过上述代码,训练得到的结果如下图所示,


请添加图片描述

结果虽然不是很好,但是我觉得已经很不多了,在测试集上的准确率差不多达到0.7了。为了节省计算资源,我把图片压缩成32*32,连我们人眼都很难分辨出哪个是蚂蚁,哪个是蜜蜂。另外,我这个模型是完全从0开始训练的,隔壁在预训练模型的基础上进行训练得到的效果好像没好多少。。。

你可能感兴趣的:(我的实践:通过蚂蚁、蜜蜂二分类问题了解如何基于Pytorch构建分类模型)