Pytorch版的Efficientnet訓練自己的資料集

文章目錄

• 前言
• 一、環境搭建
• 二、資料準備
• • 1.資料擺放
  • 2、訓練集和驗證集切分
• 三、訓練
• • 1.預訓練模型下載下傳
  • 2.加載模型
  • 3.資料讀取部分
  • 4、學習率衰減政策
  • 5、完整訓練代碼：
• 四、測試
• • 1、完整測試代碼：
  • 2、結果：
• 總結

前言

  最近，自己需要一個分類網絡來完成一項任務，于是便想起了身邊人推薦過的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是類别映射表，如我的：

這個順序得記住，在後邊實際測試的時候會用到，也可以自己加點代碼将它寫入到檔案中。

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

總結

  以上就是本篇的内容，代碼中可以根據情況可更改的參數：圖像尺寸、資料集或者模型的路徑、優化器的選取及學習率的設定等等，以上參數都是針對我的電腦來設定的。

  代碼中有的地方可能自己有點誤解或者寫錯的地方，望各位大佬指正一下，謝謝。