【代码】训练一个CNN识别小动物

好啦，前面我们学习了CNN的核心思想，那么这节我们就设计一个CNN来识别图片，之前一直都是玩手写数字识别，太无聊了，这一节我们就用CNN就来识别一下小动物的类别。

一、CNN实现宠物识别-数据准备

1.1. CNN实现宠物识别-原始数据

首先，我们得找点数据，这里我们就以https://www.robots.ox.ac.uk/~vgg/data/pets/上的宠物数据为例。

宠物数据共有7349个样本，其中12种猫样本共2371个，25种狗样本共4978个，如下：

共计37种细分类别宠物样本，平均每个细分类别大概200个样本左右。

1.2. CNN实现宠物识别-建模样本

为简化问题，我们将数据下载后，随便挑出三个细分类别的宠物，每种宠物各200个样本，然后在每种宠物选择190个作为训练样本，10个作为验证样本，并分别打上标签。处理后的数据样本如下所示：

如图所示，将3种宠物570个训练样本放到train_img文件夹，30个验证样本放到val_img文件夹，并将训练样本图片名称与类别标签{0,1,2}存放到train_label.csv文件，对应地，验证样本的标签存放在val_label.csv文件。

笔者将所有文件存放在D:\pytorch\imgdata\pet下，大家可以根据自己的情况放在自己喜欢的地方。

处理后的文件下载地址：《》

二、CNN实现宠物识别-模型设计

2.1. 宠物识别-数据处理

由于宠物图片的输入大小是不固定的，几百×几百左右，基本为三通道的RGB图像，有个别为四通道的RGBA图像，为了方便模型的输入统一尺寸，这里笔者粗略地将其统一改为256×256的统一尺寸，并统一只截取前三个通道。

2.2. 宠物识别-模型设计

根据CNN的设计思想，我们需要将256×256的输入使用"卷积+池化"进行层层压缩，使特征个数压缩到全连接层的承受范围内，整体设计思路如下：

如图所示，先用4个带池化的卷积层先将FeatureMap压缩到8×8，此时FeatureMap相对较小，再用一个卷积将FeatureMap一次性压缩为1×1即可，同时，在整个过程中，将通道逐步提升为1024，最后通过两个全连接层来拟合输出，最终输出时经过softmax，就可以得到各个类别的概率。

模型详细配置如下：

• 笔者对CNN的设计核心思想

我在设计的过程中，主要紧靠以下几点要求：

1. 卷积部分的输出要压缩到"全连接层能接受的输入个数"。

2. 卷积部分压缩特征的过程则遵循“FeatureMap在减小，通道在增大”。

3. 为了避免参数过多，在通道维度上尽量避免输入输出通道同时过大。

事实上，在整个设计过程中，我是比较粗糙、随意的，毕竟只是一个用于学习的例子，没有追求极致的必要。

三、CNN实现宠物识别-代码实现

好了，有了上面的设计，下面我们就直接用pytorch来实现就可以了。

具体代码如下：

# 本代码用于训练一个CNN模型来识别宠物类别
# 本代码来自《老饼讲解-深度学习》www.bbblearn.com
import os
import numpy  as np
import pandas as pd
import torch
from   torch import nn
from   torch.utils.data   import DataLoader
import torchvision
from torchvision.io   import read_image
from torch.utils.data import Dataset

#--------------------模型结构---------------------------
# 卷积神经网络的结构
class ConvNet(nn.Module):
    def __init__(self,in_channel,num_classes):
        super(ConvNet, self).__init__()
        self.nn_stack=nn.Sequential(
            #--------------C1层-------------------
            nn.Conv2d(in_channel,16, kernel_size=11,stride=2,padding=5),
            nn.ReLU(inplace=True),  
            nn.AvgPool2d(kernel_size=2,stride=2),
            # 输出64*64
            #--------------C2层-------------------
            nn.Conv2d(16,32, kernel_size=5,stride=1,padding=2),
            nn.ReLU(inplace=True),
            nn.AvgPool2d(kernel_size=2,stride=2),
            # 输出32*32
            #--------------C3层-------------------
            nn.Conv2d(32,64, kernel_size=5,stride=1,padding=2),
            nn.ReLU(inplace=True),
            nn.AvgPool2d(kernel_size=2,stride=2),
            # 输出16*16
            #--------------C4层-------------------
            nn.Conv2d(64,128, kernel_size=5,stride=1,padding=2),
            nn.ReLU(inplace=True),
            nn.AvgPool2d(kernel_size=2,stride=2),
            # 输出8*8
            #--------------C5层-------------------
            nn.Conv2d(128,1024,kernel_size=8,stride=1,padding=0),
            # 输出1*1*1*1024
            #--------------全连接层F6----------
            nn.Flatten(),       # 对C5的结果进行展平
            nn.Linear(1024, 256),  
            nn.ReLU(inplace=True),      
            nn.Dropout(p=0.5),                                      
            #--------------全连接层F7----------                                  
            nn.Linear(256, num_classes)                                          
            )                                                                    
    def forward(self, x):                                                        
        p = self.nn_stack(x)                                                     
        return p                                                                 
																				 
#--------------------数据处理---------------------------                         
# 自定义DataSet数据类                                                            
class CustomImageDataset(Dataset):                                               
    def __init__(self, annotations_file, img_dir, transform = None):             
        self.img_labels = pd.read_csv(annotations_file, header = None)            # 从CSV中读取图象标签
        self.img_dir    = img_dir                                                 # 存放图片的文件夹
        self.transform  = transform                                               # 图片的转换函数
                                                                                  
    def __len__(self):                                                            
        return len(self.img_labels)                                               # 标签的长度就是样本个数
                                                                                  
    def __getitem__(self, idx):                                                   
        img_path = os.path.join(self.img_dir, self.img_labels.iloc[idx, 0])       # 图片路径
        image    = read_image(img_path).float()                                   # 读取图片
        image    = image[0:3,:,:]                                                 # 避免4通道图片
        label    = self.img_labels.iloc[idx, 1]                                   # 读取标签
        if self.transform:                                                        # 如果有图片的转换函数
            image = self.transform(image)                                         # 就对图片进行转换
        return image, label                                                       # 返回图片和标签    
																				 
#-----------------------模型训练-------------------------                        
# 参数初始化函数                                                                 
def init_param(model):                                                           
    # 初始化权重阈值                                                              
    param_list = list(model.named_parameters())                                   # 将模型的参数提取为列表
    for i in range(len(param_list)):                                              # 逐个初始化权重、阈值
        is_weight = i%2==0                                                        # 偶数层是权重参数,奇数是阈值
        if is_weight:                                                             
            torch.nn.init.normal_(param_list[i][1],mean=0,std=0.01)               # 权重层以N(0,0.01)随机初始化
        else:                                                                     
           torch.nn.init.constant_(param_list[i][1],val=0)                        # 阈值初始化为0
                                                                                     
# 训练函数                                                                           
def train(dataloader,valloader,model,optimizer,epochs,goal,device):                            
    for epoch in range(epochs):                                                      
        err_num  = 0                                                              # 本次epoch评估错误的样本
        eval_num = 0                                                              # 本次epoch已评估的样本
        print('-----------当前epoch:',str(epoch),'----------------')              
        for batch, (imgs, labels) in enumerate(dataloader):                       
		    # -----训练模型-----                                                  
            x, y = imgs.to(device), labels.to(device)                             # 将数据发送到设备
            optimizer.zero_grad()                                                 # 将优化器里的参数梯度清空
            py   = model(x)                                                       # 计算模型的预测值   
            loss = lossFun(py, y)                                                 # 计算损失函数值
            loss.backward()                                                       # 更新参数的梯度
            optimizer.step()                                                      # 更新参数
			# ----计算错误率----                                                  
            idx      = torch.argmax(py,axis=1)                                    # 模型的预测类别
            eval_num = eval_num + len(idx)                                        # 更新本次epoch已评估样本
            err_num  = err_num +sum(y != idx)                                     # 更新本次epoch评估错误样本
            if(batch%10==0):                                                      # 每10批打印一次结果
                print('err_rate:',err_num/eval_num)                               # 打印错误率
        # -----------验证数据误差---------------------------                      
        model.eval()                                                              # 将模型调整为评估状态
        val_acc_rate = calAcc(model,valLoader,device)                             # 计算验证数据集的准确率
        model.train()                                                             # 将模型调整回训练状态
        print("验证数据的准确率:",val_acc_rate)                                   # 打印准确率    
        if((err_num/eval_num)<=goal):                                             # 检查退出条件
            break                                                                 
    print('训练步数',str(epoch),',最终训练误差',str(err_num/eval_num))            
                                                                                  
# 计算数据集的准确率                                                              
def calAcc(model,dataLoader,device):                                              
    py = np.empty(0)                                                              # 初始化预测结果
    y  = np.empty(0)                                                              # 初始化真实结果
    for batch, (imgs, labels) in enumerate(dataLoader):                           # 逐批预测
        cur_py =  model(imgs.to(device))                                          # 计算网络的输出
        cur_py = torch.argmax(cur_py,axis=1)                                      # 将最大者作为预测结果
        py     = np.hstack((py,cur_py.detach().cpu().numpy()))                    # 记录本批预测的y
        y      = np.hstack((y,labels))                                            # 记录本批真实的y
    acc_rate = sum(y==py)/len(y)                                                  # 计算测试样本的准确率
    return acc_rate                                                              
																				 
#---------------主流程脚本-----------------------                     
#---------加载数据-----------                                                    
trans = torchvision.transforms.Compose([                                         
    torchvision.transforms.Resize((256,256),antialias=True) ,                     # 改变图片大小
])                                                                                
# 通过数据类读取训练数据                                                          
train_img_dir    = "D:\\pytorch\\imgdata\\pet\\train_img"                         # 训练数据-图象文件夹
train_label_file = "D:\\pytorch\\imgdata\\pet\\train_label.csv"                   # 训练数据-标签文件
train_data = CustomImageDataset(train_label_file,train_img_dir,transform=trans)   # 初始化训练数据类
# 通过数据类读取验证数据                                                               
val_img_dir    = "D:\\pytorch\\imgdata\\pet\\train_img"                           # 验证数据-图象文件夹
val_label_file = "D:\\pytorch\\imgdata\\pet\\train_label.csv"                     # 验证数据-标签文件
val_data = CustomImageDataset(val_label_file,val_img_dir,transform=trans)         # 初始化验证数据类    
																				 
#-------模型训练--------                                                         
trainLoader = DataLoader(train_data, batch_size=100, shuffle=True)                # 训练数据装载到DataLoader
valLoader   = DataLoader(val_data  , batch_size=100)                              # 验证数据装载到DataLoader 
device      = torch.device('cuda' if torch.cuda.is_available() else 'cpu')        # 设置训练设备  
model       = ConvNet(in_channel=3,num_classes=3).to(device)                      # 初始化模型,并发送到设备  
init_param(model)                                                                 
lossFun     = torch.nn.CrossEntropyLoss()                                         # 损失函数为交叉熵损失函数
optimizer   = torch.optim.SGD(model.parameters(), lr=0.001,momentum =0.9)         # 初始化优化器
train(trainLoader,valLoader,model,optimizer,1000,0.01,device)                     # 训练模型
																				 
# -----------模型效果评估---------------------------                             
model.eval()                                                                      # 切到评估状态(屏蔽Dropout)
train_acc_rate = calAcc(model,trainLoader,device)                                 # 计算训练数据集的准确率
print("训练数据的准确率:",train_acc_rate)                                         # 打印准确率
val_acc_rate = calAcc(model,valLoader,device)                                     # 计算验证数据集的准确率
print("验证数据的准确率:",val_acc_rate)                                           # 打印准确率

运行结果如下：

可以看到，模型的训练准确率与验证数据的准确率都达到了100%，说明模型是有效的。

结束语

这里呢，我们只挑选了3个类别，是为了简化问题，降低学习成本，而它的输入样本为256*256*3的图片，则是为了体现CNN在解决图像识别问题上的独特优势，同时，使用宠物数据，也可以让大家熟悉一下如何导入自己的图片数据。从本例就可以看到，随便设计一个CNN，都能非常好地解决一个不太难的图像类别识别问题，比MLP好用多了。