Pytorch版的Efficientnet训练自己的数据集

---

前言

  最近，自己需要一个分类网络来完成一项任务，于是便想起了身边人推荐过的Efficientnet，据说效果是较为稳定的，所以自己来一探究竟，示例的话就用个最简单的二分类吧。

---

一、环境搭建

本人使用的环境为:

python3.6
torch=1.5
torchvision =0.6.0
opencv-python=4.5.1.48

以上这些仅供参考，无需一致，重要的使我们还需要安装pytorch集合进来的Efficientnet模块，在我们要使用的python环境下,执行命令

pip install efficientnet_pytorch

其他依赖项到时逐个安装即可。

---

二、数据准备

1.数据摆放

  原始数据摆放如下：

  也就是以类别名来命名文件夹名，将对应的类别图片放置对应的文件夹下，一般来说，分类任务的数据集大多都是这样来摆放的。

2、训练集和验证集切分

  这一步只需要运行dataset.py即可，它会按照我们制定的比例将我们的数据集进行切分开，同时，为了减少直接resize带来的图片变形的弊端，这里在切分的同时我对数据还进行补边的操作，也就是将数据尽量变为正方形的样子，代码如下

#为efficientnet训练分类的数据进行预处理（训练集切分+补边）
import os
import glob
import cv2
import random
from pathlib import Path


#补边,这一步主要是为了将图片填充为正方形，防止直接resize导致图片变形
def expend_img(img):
    '''
    :param img: 图片数据
    :return:
    '''
    fill_pix=[122,122,122] #填充色素，可自己设定
    h,w=img.shape[:2]
    if h>=w: #左右填充
        padd_width=int(h-w)//2
        padd_top,padd_bottom,padd_left,padd_right=0,0,padd_width,padd_width #各个方向的填充像素
    elif h<w: #上下填充
        padd_high=int(w-h)//2
        padd_top,padd_bottom,padd_left,padd_right=padd_high,padd_high,0,0 #各个方向的填充像素
    new_img = cv2.copyMakeBorder(img,padd_top,padd_bottom,padd_left,padd_right,cv2.BORDER_CONSTANT, value=fill_pix)
    return new_img


#切分训练集和测试集，并进行补边处理
def split_train_test(img_dir,save_dir,train_val_num):
    '''
    :param img_dir: 原始图片路径，注意是所有类别所在文件夹的上一级目录
    :param save_dir: 保存图片路径
    :param train_val_num: 切分比例
    :return:
    '''
    img_dir_list=glob.glob(img_dir+os.sep+"*")#获取每个类别所在的路径（一个类别对应一个文件夹）
    for class_dir in img_dir_list:
        class_name=class_dir.split(os.sep)[-1] #获取当前类别
        img_list=glob.glob(class_dir+os.sep+"*") #获取每个类别文件夹下的所有图片
        all_num=len(img_list) #获取总个数
        train_list=random.sample(img_list,int(all_num*train_val_num)) #训练集图片所在路径
        save_train=save_dir+os.sep+"train"+os.sep+class_name
        save_val=save_dir+os.sep+"val"+os.sep+class_name
        os.makedirs(save_train,exist_ok=True)
        os.makedirs(save_val,exist_ok=True) #建立对应的文件夹
        print(class_name+" trian num",len(train_list))
        print(class_name+" val num",all_num-len(train_list))
        #保存切分好的数据集
        for imgpath in img_list:
            imgname=Path(imgpath).name #获取文件名
            if imgpath in train_list:
                img=cv2.imread(imgpath)
                new_img=expend_img(img)
                cv2.imwrite(save_train+os.sep+imgname,new_img)
            else: #将除了训练集意外的数据均视为验证集
                img = cv2.imread(imgpath)
                new_img = expend_img(img)
                cv2.imwrite(save_val + os.sep + imgname, new_img)

    print("split train and val finished !")

这里，也对代码内容和相关参数进行了注释，理解起来应该不是很难
运行它的时候，我们只需要调用split_train_test（）函数，输入指定的参数（3个）即可，需要注意的是，这里给的原始图片的路径是所有类别文件夹的上一级，程序会依次遍历它下面的各个文件夹来进行切分运行完成后，会生成对应的训练集和测试集，如下图：

train和val里也会有生成各个类别的文件夹用于储存不同类别的数据，需要注意的是，这里存放的数据是我经过补边之后的，对原路径的数据集不会有改动，填充颜色我默认设置为了灰色，可以根据自己爱好在代码中自行更改

三、训练

1.预训练模型下载

这里可以对代码进行更改，使它自动下载模型，我是觉得慢，所以手动下载了，网址如下：

'''
efficientnet-b0: https://github.com/lukemelas/EfficientNet-PyTorch/releases/download/1.0/efficientnet-b0-355c32eb.pth
efficientnet-b1: https://github.com/lukemelas/EfficientNet-PyTorch/releases/download/1.0/efficientnet-b1-f1951068.pth
efficientnet-b2: https://github.com/lukemelas/EfficientNet-PyTorch/releases/download/1.0/efficientnet-b2-8bb594d6.pth
efficientnet-b3: https://github.com/lukemelas/EfficientNet-PyTorch/releases/download/1.0/efficientnet-b3-5fb5a3c3.pth
efficientnet-b4: https://github.com/lukemelas/EfficientNet-PyTorch/releases/download/1.0/efficientnet-b4-6ed6700e.pth
efficientnet-b5: https://github.com/lukemelas/EfficientNet-PyTorch/releases/download/1.0/efficientnet-b5-b6417697.pth
efficientnet-b6: https://github.com/lukemelas/EfficientNet-PyTorch/releases/download/1.0/efficientnet-b6-c76e70fd.pth
efficientnet-b7: https://github.com/lukemelas/EfficientNet-PyTorch/releases/download/1.0/efficientnet-b7-dcc49843.pth
'''

我选用的是b0

2.加载模型

代码如下（示例）：

        base_model = EfficientNet.from_name('efficientnet-b0') #加载模型，使用b几的就改为b几
        state_dict = torch.load(self.weights)
        base_model.load_state_dict(state_dict)
        # 修改全连接层
        num_ftrs = base_model._fc.in_features
        base_model._fc = nn.Linear(num_ftrs, self.class_num)
        self.model = base_model.to(device)

3.数据读取部分

这里对数据进行了指定的数据变换（增强），可以根据需求进行删改，代码如下：

	#数据处理
    def process(self):
        # 数据增强
        data_transforms = {
            'train': transforms.Compose([
                transforms.Resize((self.imgsz, self.imgsz)),  # resize
                transforms.CenterCrop((self.imgsz, self.imgsz)),  # 中心裁剪
                transforms.RandomRotation(10),  # 随机旋转，旋转范围为【-10,10】
                transforms.RandomHorizontalFlip(p=0.2),  # 水平镜像
                transforms.ToTensor(),  # 转换为张量
                transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])  # 标准化
            ]),
            "val": transforms.Compose([
                transforms.Resize((self.imgsz, self.imgsz)),  # resize
                transforms.CenterCrop((self.imgsz, self.imgsz)),  # 中心裁剪
                transforms.ToTensor(),  # 张量转换
                transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
            ])
        }

        # 定义图像生成器
        image_datasets = {x: datasets.ImageFolder(os.path.join(self.img_dir, x), data_transforms[x]) for x in
                          ['train', 'val']}
        # 得到训练集和验证集
        trainx = DataLoader(image_datasets["train"], batch_size=self.batch_size, shuffle=True, drop_last=True)
        valx = DataLoader(image_datasets["val"], batch_size=self.batch_size, shuffle=True, drop_last=True)

        b = image_datasets["train"].class_to_idx  # id和类别对应

        return trainx,valx,b

ImageFolder()这个函数，如果有人不清楚的，可以进行百度，返回的b是类别映射表，如我的：

{'cat': 0, 'dog': 1}

这个顺序得记住，在后边实际测试的时候会用到，也可以自己加点代码将它写入到文件中。

4、学习率衰减策略

这里的方案是先从初始值上升，然后在保持不动，然后在进行指数衰减，代码如下：

    # 学习率慢热加下降
    def lrfn(self,num_epoch, optimzer):
        lr_start = 0.00001  # 初始值
        max_lr = 0.0004  # 最大值
        lr_up_epoch = 10  # 学习率上升10个epoch
        lr_sustain_epoch = 5  # 学习率保持不变
        lr_exp = .8  # 衰减因子
        if num_epoch < lr_up_epoch:  # 0-10个epoch学习率线性增加
            lr = (max_lr - lr_start) / lr_up_epoch * num_epoch + lr_start
        elif num_epoch < lr_up_epoch + lr_sustain_epoch:  # 学习率保持不变
            lr = max_lr
        else:  # 指数下降
            lr = (max_lr - lr_start) * lr_exp ** (num_epoch - lr_up_epoch - lr_sustain_epoch) + lr_start
        for param_group in optimzer.param_groups:
            param_group['lr'] = lr
        return optimzer

其中，参数 lr_sustain_epoch、max_lr、lr_up_epoch、 lr_sustain_epoch等均可以按照需求进行调整，非固定值

5、完整训练代码：

from torchvision import datasets,transforms
import torch
import torch.optim as optim
import torch.nn as nn
from torch.utils.data import DataLoader
from efficientnet_pytorch import EfficientNet
import os
import time
import argparse

device="cuda" if torch.cuda.is_available() else "cpu"

class Efficientnet_train():
    def __init__(self,opt):
        self.epochs=opt.epochs #训练周期
        self.batch_size=opt.batch_size #batch_size
        self.class_num=opt.class_num #类别数
        self.imgsz=opt.imgsz #图片尺寸
        self.img_dir=opt.img_dir #图片路径
        self.weights=opt.weights #模型路径
        self.save_dir=opt.save_dir #保存模型路径
        self.lr=opt.lr #初始化学习率
        self.moment=opt.m #动量
        base_model = EfficientNet.from_name('efficientnet-b0') #记载模型，使用b几的就改为b几
        state_dict = torch.load(self.weights)
        base_model.load_state_dict(state_dict)
        # 修改全连接层
        num_ftrs = base_model._fc.in_features
        base_model._fc = nn.Linear(num_ftrs, self.class_num)
        self.model = base_model.to(device)
        # 交叉熵损失函数
        self.cross = nn.CrossEntropyLoss()
        # 优化器
        self.optimzer = optim.SGD((self.model.parameters()), lr=self.lr, momentum=self.moment, weight_decay=0.0004)

        #获取处理后的数据集和类别映射表
        self.trainx,self.valx,self.b=self.process()
        print(self.b)
    def __call__(self):
        best_acc = 0
        self.model.train(True)
        for ech in range(self.epochs):
            optimzer1 = self.lrfn(ech, self.optimzer)

            print("----------Start Train Epoch %d----------" % (ech + 1))
            # 开始训练
            run_loss = 0.0  # 损失
            run_correct = 0.0  # 准确率
            count = 0.0  # 分类正确的个数

            for i, data in enumerate(self.trainx):

                inputs, label = data
                inputs, label = inputs.to(device), label.to(device)

                # 训练
                optimzer1.zero_grad()
                output = self.model(inputs)

                loss = self.cross(output, label)
                loss.backward()
                optimzer1.step()

                run_loss += loss.item()  # 损失累加
                _, pred = torch.max(output.data, 1)
                count += label.size(0)  # 求总共的训练个数
                run_correct += pred.eq(label.data).cpu().sum()  # 截止当前预测正确的个数
                #每隔100个batch打印一次信息，这里打印的ACC是当前预测正确的个数/当前训练过的的个数
                if (i+1)%100==0:
                    print('[Epoch:{}__iter:{}/{}] | Acc:{}'.format(ech + 1,i+1,len(self.trainx), run_correct/count))

            train_acc = run_correct / count
            # 每次训完一批打印一次信息
            print('Epoch:{} | Loss:{} | Acc:{}'.format(ech + 1, run_loss / len(self.trainx), train_acc))

            # 训完一批次后进行验证
            print("----------Waiting Test Epoch {}----------".format(ech + 1))
            with torch.no_grad():
                correct = 0.  # 预测正确的个数
                total = 0.  # 总个数
                for inputs, labels in self.valx:
                    inputs, labels = inputs.to(device), labels.to(device)
                    outputs = self.model(inputs)

                    # 获取最高分的那个类的索引
                    _, pred = torch.max(outputs.data, 1)
                    total += labels.size(0)
                    correct += pred.eq(labels).cpu().sum()
                test_acc = correct / total
                print("批次%d的验证集准确率" % (ech + 1), correct / total)
            if best_acc < test_acc:
                best_acc = test_acc
                start_time=(time.strftime("%m%d",time.localtime()))
                save_weight=self.save_dir+os.sep+start_time #保存路径
                os.makedirs(save_weight,exist_ok=True)
                torch.save(self.model, save_weight + os.sep + "best.pth")

    #数据处理
    def process(self):
        # 数据增强
        data_transforms = {
            'train': transforms.Compose([
                transforms.Resize((self.imgsz, self.imgsz)),  # resize
                transforms.CenterCrop((self.imgsz, self.imgsz)),  # 中心裁剪
                transforms.RandomRotation(10),  # 随机旋转，旋转范围为【-10,10】
                transforms.RandomHorizontalFlip(p=0.2),  # 水平镜像
                transforms.ToTensor(),  # 转换为张量
                transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])  # 标准化
            ]),
            "val": transforms.Compose([
                transforms.Resize((self.imgsz, self.imgsz)),  # resize
                transforms.CenterCrop((self.imgsz, self.imgsz)),  # 中心裁剪
                transforms.ToTensor(),  # 张量转换
                transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
            ])
        }

        # 定义图像生成器
        image_datasets = {x: datasets.ImageFolder(os.path.join(self.img_dir, x), data_transforms[x]) for x in
                          ['train', 'val']}
        # 得到训练集和验证集
        trainx = DataLoader(image_datasets["train"], batch_size=self.batch_size, shuffle=True, drop_last=True)
        valx = DataLoader(image_datasets["val"], batch_size=self.batch_size, shuffle=True, drop_last=True)

        b = image_datasets["train"].class_to_idx  # id和类别对应

        return trainx,valx,b


    # 学习率慢热加下降
    def lrfn(self,num_epoch, optimzer):
        lr_start = 0.00001  # 初始值
        max_lr = 0.0004  # 最大值
        lr_up_epoch = 10  # 学习率上升10个epoch
        lr_sustain_epoch = 5  # 学习率保持不变
        lr_exp = .8  # 衰减因子
        if num_epoch < lr_up_epoch:  # 0-10个epoch学习率线性增加
            lr = (max_lr - lr_start) / lr_up_epoch * num_epoch + lr_start
        elif num_epoch < lr_up_epoch + lr_sustain_epoch:  # 学习率保持不变
            lr = max_lr
        else:  # 指数下降
            lr = (max_lr - lr_start) * lr_exp ** (num_epoch - lr_up_epoch - lr_sustain_epoch) + lr_start
        for param_group in optimzer.param_groups:
            param_group['lr'] = lr
        return optimzer
#参数设置
def parse_opt():
    parser=argparse.ArgumentParser()
    parser.add_argument("--weights",type=str,default="./models/efficientnet-b0-355c32eb.pth",help='initial weights path')#预训练模型路径
    parser.add_argument("--img-dir",type=str,default="",help="train image path") #数据集的路径
    parser.add_argument("--imgsz",type=int,default=224,help="image size") #图像尺寸
    parser.add_argument("--epochs",type=int,default=50,help="train epochs")#训练批次
    parser.add_argument("--batch-size",type=int,default=4,help="train batch-size") #batch-size
    parser.add_argument("--class_num",type=int,default=2,help="class num") #类别数
    parser.add_argument("--lr",type=float,default=0.0001,help="Init lr") #学习率初始值
    parser.add_argument("--m",type=float,default=0.9,help="optimer momentum") #动量
    parser.add_argument("--save-dir",type=str,default="./weight",help="save models dir")#保存模型路径
    opt=parser.parse_known_args()[0]
    return opt

if __name__ == '__main__':
    opt=parse_opt()
    models=Efficientnet_train(opt)
    models()

只需要将对应的参数设置为自己的就可

---

四、测试

1、完整测试代码：

这里，话不多说，直接上代码

import torch
import os
import torchvision
import glob
from PIL import Image
import cv2
import argparse
device="cuda" if torch.cuda.is_available() else "cpu"
#参数设置
def parser_opt():
    parser=argparse.ArgumentParser()
    parser.add_argument("--test-dir",type=str,default=r"")
    parser.add_argument("--weights",type=str,default="",help="model path")
    parser.add_argument("--imgsz",type=int,default=224,help="test image size")
    opt=parser.parse_known_args()[0]
    return opt
#测试图片
class Test_model():
    def __init__(self,opt):
        self.imgsz=opt.imgsz #测试图片尺寸
        self.img_dir=opt.test_dir #测试图片路径

        self.model=(torch.load(opt.weights)).to(device) #加载模型
        self.model.eval()
        self.class_name=[] #类别信息
    def __call__(self):
        #图像转换
        data_transorform=torchvision.transforms.Compose([
                torchvision.transforms.Resize((224,224)),
                torchvision.transforms.CenterCrop((224,224)),
                torchvision.transforms.ToTensor(),
                torchvision.transforms.Normalize(mean=[0.485,0.456,0.406],std=[0.229,0.224,0.225])
            ])
        img_list=glob.glob(self.img_dir+os.sep+"*.jpg")
      
        for imgpath in img_list:
            img=cv2.imread(imgpath)
            new_img=self.expend_img(img) #补边
            img=Image.fromarray(new_img)
            img=data_transorform(img) #转换
            img=torch.reshape(img,(-1,3,self.imgsz,self.imgsz)).to(device) #维度转换[B,C,H,W]
            pred=self.model(img)
            _,pred=torch.max(pred,1)
            outputs = self.class_name[pred]
            print("Image path:",imgpath," pred:",outputs)

    #补边为正方形
    def expend_img(self,img,fill_pix=122):
        '''
        :param img: 图片数据
        :param fill_pix: 填充像素，默认为灰色，自行更改
        :return:
        '''
        h,w=img.shape[:2] #获取图像的宽高
        if h>=w: #左右填充
            padd_width=int(h-w)//2
            padd_h,padd_b,padd_l,padd_r=0,0,padd_width,padd_width #获取上下左右四个方向需要填充的像素

        elif h<w: #上下填充
            padd_high=int(w-h)//2
            padd_h,padd_b,padd_l,padd_r=padd_high,padd_high,0,0

        new_img = cv2.copyMakeBorder(img, padd_h, padd_b, padd_l, padd_r, borderType=cv2.BORDER_CONSTANT,
                                     value=[fill_pix,fill_pix,fill_pix])
        return new_img

if __name__ == '__main__':
    opt=parser_opt()
    test_img=Test_model(opt)
    test_img()

这里，依然对测试数据进行了补边处理，相关参数在parser_opt()里给定即可
注意：self.class_name=[] 里的类别信息写为自己的，也就是3.3里让记录的那个顺序。

2、结果：

这是我测试的某类别的100张图片，图片均为新图，效果上看还可以，训练50批次

总结

  以上就是本篇的内容，代码中可以根据情况可更改的参数：图像尺寸、数据集或者模型的路径、优化器的选取及学习率的设置等等，以上参数都是针对我的电脑来设置的。

  代码中有的地方可能自己有点误解或者写错的地方，望各位大佬指正一下，谢谢。