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

---

前言

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

---

一、环境搭建

本人使用的环境为:

python3.6
torch=1.5
torchvision =0.6.0
opencv-python=4.5.1.48
1
2
3
4

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

pip install efficientnet_pytorch
1

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

---

二、数据准备

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 !")

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60

这里，也对代码内容和相关参数进行了注释，理解起来应该不是很难
运行它的时候，我们只需要调用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
'''
1
2
3
4
5
6
7
8
9
10

我选用的是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)
1
2
3
4
5
6
7

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
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30

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

{'cat': 0, 'dog': 1}
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
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16

其中，参数 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
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164

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

---

四、测试

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()
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70

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

2、结果：

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

总结

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

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