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批次

总结

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

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

  • 14
    点赞
  • 71
    收藏
    觉得还不错? 一键收藏
  • 22
    评论
### 回答1: 要使用EfficientNet训练自己的数据,您需要遵循以下步骤: 1. 准备数据集:收集和准备您的数据集,确保它们符合您的需求和标准。 2. 数据预处理:对数据进行预处理,包括图像增强、数据增强、数据划分等。 3. 构建模型:使用EfficientNet模型作为基础模型,根据您的需求和数据集进行微调和修改。 4. 训练模型:使用训练数据集对模型进行训练,并使用验证数据集进行验证和调整。 5. 测试模型:使用测试数据集对模型进行测试和评估。 6. 模型优化:根据测试结果对模型进行优化和改进,以提高模型的准确性和效率。 以上是训练自己数据集的一般步骤,具体实现需要根据您的具体情况进行调整和修改。 ### 回答2: EfficientNet是一种高效的神经网络架构,可用于训练自己的数据。它采用了一系列的网络结构扩展方式,通过提高网络的深度、宽度和分辨率来提升性能。 首先,为了训练自己的数据集,我们需要准备数据。将数据集划分为训练集和验证集,并对数据进行标注和预处理,以便让模型能够更好地学习。预处理可能包括图像增强和数据扩充等操作,以增加数据的多样性。 接下来,我们需要下载预训练的EfficientNet模型权重。这些权重可以作为预训练模型的初始参数,迁移到我们的数据集上进行微调。这样可以在有限的数据集上获得更好的性能。 然后,将自己的数据集输入EfficientNet模型进行训练。使用合适的优化算法如随机梯度下降(SGD)或自适应矩估计(Adam),根据训练集的样本进行反向传播和参数更新。同时,使用验证集评估模型的性能,避免过拟合。 在训练过程中,可以使用学习率衰减策略来调整学习率,以提高训练的稳定性和效果。此外,还可以在训练过程中使用正则化技术,如Dropout或L2正则化,以防止过拟合。 当训练达到预定的停止条件时,可以保存模型并使用测试集进行性能评估。测试集应与训练集和验证集相互独立,以确保模型的泛化能力。 最后,根据模型在测试集上的性能,可以进行模型调整、超参数调优等进一步优化工作,以获得更好的结果。 总之,使用EfficientNet训练自己的数据需要准备数据集、下载预训练权重、迁移到自己的数据集上进行微调、选择优化算法、进行训练和验证、调整超参数等步骤。通过这些步骤,就可以以高效的方式训练自己的数据。 ### 回答3: 要训练自己的数据集,可以使用EfficientNet这个强大的模型。首先,需要确保数据集已经准备好,并且正确地标注了每个样本。接下来,根据数据集的大小和复杂度,选择合适的EfficientNet模型本,如EfficientNet-B0到EfficientNet-B7。 然后,根据数据集的情况,调整EfficientNet模型的超参数和训练参数。例如,可以调整学习率、批大小和训练轮数等参数,以获得更好的训练效果。 在训练过程中,可以使用一些常见的数据增强技术,如旋转、平移、裁剪和翻转等,来增加数据集的多样性和数量。这样可以提高模型的泛化能力和稳定性。 为了训练EfficientNet模型,可以使用流行的深度学习框架,如TensorFlow或PyTorch。使用这些框架可以方便地定义模型结构、加载数据集和进行模型训练。 在训练过程中,可以监控模型的训练损失和验证准确率等指标,以便了解模型的训练进展。如果发现模型出现过拟合或欠拟合等问题,可以尝试调整超参数或增加更多的训练数据来改善模型效果。 当模型训练完成后,可以使用测试集对模型进行评估,并计算模型的准确率、精确率、召回率和F1分数等指标。根据评估结果,可以进一步调整模型和训练策略,以获得更好的性能。 最后,可以使用训练好的EfficientNet模型对新的未知样本进行预测。将样本输入到模型中,即可得到相应的分类结果。 总而言之,EfficientNet是一个强大的模型,可以通过调整超参数和训练策略来训练自己的数据集。通过不断优化和改进,可以获得准确性高并且泛化能力强的模型。

“相关推荐”对你有帮助么?

  • 非常没帮助
  • 没帮助
  • 一般
  • 有帮助
  • 非常有帮助
提交
评论 22
添加红包

请填写红包祝福语或标题

红包个数最小为10个

红包金额最低5元

当前余额3.43前往充值 >
需支付:10.00
成就一亿技术人!
领取后你会自动成为博主和红包主的粉丝 规则
hope_wisdom
发出的红包
实付
使用余额支付
点击重新获取
扫码支付
钱包余额 0

抵扣说明:

1.余额是钱包充值的虚拟货币,按照1:1的比例进行支付金额的抵扣。
2.余额无法直接购买下载,可以购买VIP、付费专栏及课程。

余额充值