pytorch 训练EfficientnetV2

---

前言

  前不久用pytorch复现了efficientnetv2的网络结构，但是后边自己一直有其他事情再做，所以训练部分的文章拖到了现在。关于Efficientnet的部分文章链接可参考如下：
EfficientnetV1训练
Flask部署EfficientnetV1分类网络
pytorch构建EfficientnetV2网络结构
然后，本篇的训练代码也是基于v1的训练代码完成，所以和官方会有差距。

---

一、数据摆放

  训练数据摆放方式可以参考EfficientnetV1训练的文章，大致如下：

在train和val文件夹下又有各个类别的数据，这里拿猫狗分类来说，如下：

二、训练

  这里直接给了训练代码，部分注释已在代码中添加，仔细看会发现和V1的训练代码差不多，只是改了封装方式和几个参数而已。需要注意的是需要将上边提到的用pytorch复现的efficientnetV2网络结构copy至model.py中，摆放如下

代码详情如下：

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

device="cuda" if torch.cuda.is_available() else "cpu"
#数据处理
def process(opt):
    # 数据增强
    data_transforms = {
    
    
        'train': transforms.Compose([
            # transforms.Resize((self.imgsz, self.imgsz)),  # resize
            transforms.CenterCrop((opt.imgsz, opt.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((opt.imgsz, opt.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(opt.img_dir, x), data_transforms[x]) for x in
                      ['train', 'val']}
    # 得到训练集和验证集
    trainx = DataLoader(image_datasets["train"], batch_size=opt.batch_size, shuffle=True, drop_last=True)
    valx = DataLoader(image_datasets["val"], batch_size=opt.batch_size, shuffle=True, drop_last=True)

    b = image_datasets["train"].class_to_idx  # id和类别对应
    print(b)
    return trainx,valx,b
#训练
def train(opt):
    start_time = (time.strftime("%m%d_%H%M", time.localtime()))
    save_weight = opt.save_dir + os.sep + start_time  # 保存路径
    os.makedirs(save_weight, exist_ok=True)
    model=EfficientnetV2(opt.model_type,opt.class_num).cuda()
    best_acc = 0
    best_epoch = 0  # 准确率最高的模型的训练周期
    model.train(True)
    # 优化器
    optimzer=optim.SGD(model.parameters(),lr=opt.lr,momentum=opt.m,weight_decay=0.0004)
    cross = nn.CrossEntropyLoss() #损失函数
    trainx, valx, b = process(opt)

    for ech in range(opt.epochs):
        optimzer1 = lrfn(ech, optimzer,opt.lr)

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

        for i, data in enumerate(trainx):

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

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

            loss = 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) % 500 == 0:
                print('[Epoch:{}__iter:{}/{}] | Acc:{}'.format(ech + 1, i + 1, len(trainx), run_correct / count))

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

        # 训完一批次后进行验证
        print("----------Waiting Test Epoch {}----------".format(ech + 1))
        with torch.no_grad():
            correct = 0.  # 预测正确的个数
            total = 0.  # 总个数
            for inputs, labels in valx:
                inputs, labels = inputs.to(device), labels.to(device)
                outputs = 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), test_acc.cpu().detach().numpy())
        if best_acc < test_acc:
            best_acc = test_acc
            best_epoch = ech + 1

            torch.save(model, save_weight + os.sep + "best.pth")
        print(f'best epoch : {
      
      best_epoch}, best accuracy : {
      
      best_acc}')

#学习率设置
def lrfn(num_epoch,optim,lr):
    lr_start=lr
    max_lr=0.01
    lr_up_epoch = 5  # 学习率上升批次
    lr_sustain_epoch = 10  # 学习率保持不变
    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 optim.param_groups:
        param_group['lr'] = lr
    return optim


def parse_opt():
    parser = argparse.ArgumentParser()
    parser.add_argument("--model_type",type=str,default="S",help="Model type") #模型选型，可选s,m,l，大小写均可
    parser.add_argument("--img-dir", type=str, default="", help="train image path")  # 数据集的路径
    parser.add_argument("--imgsz", type=int, default=480, help="image size")  # 图像尺寸
    parser.add_argument("--epochs", type=int, default=100, help="train epochs")  # 训练批次
    parser.add_argument("--batch-size", type=int, default=16, 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="",
                        help="save models dir")  # 保存模型路径
    opt = parser.parse_known_args()[0]
    return opt



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

在train.py中，只需要在parse_opt()中选取个模型类型，如S、M、L等，其次，数据集路径、类别数、保存路径等参数均设置为自己的。
注意：代码中38行附近有个print(b)，这个是类别列表，一定记住这个列表，在测试中会用，否则在测试时，类别可能都是错的。

二、测试

测试部分直接给出代码，其中第44行和第48行需要该为自己的测试数据路径和模型路径，代码如下

import torch
import torchvision
from PIL import Image
import cv2,glob,os,time
import shutil
from pathlib import Path

def expend_img(img,img_size=480,expand_pix=0):
    '''
    :param img: 图片数据
    :param fill_pix: 填充像素，默认为灰色，自行更改
    :return:
    '''
    h, w = img.shape[:2]
    if h > w and h >= img_size:  # 左右padding
        top_expand = 0
        bottom_expand = 0
        left_expand = int((h - w) / 2)
        right_expand = left_expand
        new_img = cv2.copyMakeBorder(img, top_expand, bottom_expand, left_expand, right_expand, cv2.BORDER_CONSTANT,
                                     value=expand_pix)
    elif w > h and w >= img_size:  # 上下padding
        left_expand = 0
        right_expand = 0
        top_expand = int((w - h) / 2)
        bottom_expand = top_expand
        new_img = cv2.copyMakeBorder(img, top_expand, bottom_expand, left_expand, right_expand, cv2.BORDER_CONSTANT,
                                     value=expand_pix)
    elif w < img_size and h < img_size:  # 四周padding
        left_expand = int((img_size - w) / 2)
        right_expand = left_expand
        top_expand = int((img_size - h) / 2)
        bottom_expand = top_expand
        new_img = cv2.copyMakeBorder(img, top_expand, bottom_expand, left_expand, right_expand, cv2.BORDER_CONSTANT,
                                     value=expand_pix)
    else:
        new_img = img

    new_img = cv2.resize(new_img, (img_size, img_size))
    return new_img

#模糊分类
if __name__ == '__main__':
    img_dir="" #测试数据路径
    img_list=glob.glob(img_dir+os.sep+"*.jpg")
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    #加载模型
    model=torch.load("").to(device)
    model.eval()

    class_list=[]

    for imgpath in img_list:
        img=cv2.imread(imgpath)
        s=time.time()
        img=expend_img(img)

        #PIL
        img = Image.fromarray(cv2.cvtColor(img, cv2.COLOR_BGR2RGB))  # 注意时间
        # # img= img[:, :, ::-1].transpose(2, 0, 1).copy() #注意时间

        data_transorform = torchvision.transforms.Compose([
            torchvision.transforms.ToTensor(),
            torchvision.transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
        ])
        img = data_transorform(img)
        pred_img = torch.reshape(img, (-1, 3, 480,480)).to(device)
        start_time=time.time()
        pred=model(pred_img)[0]
        pred = torch.nn.Softmax(dim=0)(pred)
        end_time=time.time()
        score, pred_id = torch.max(pred, dim=0)
        #预测类别
        pred_class=class_list[pred_id]
        e=time.time()

        print(f"{imgpath} is {pred_class},score is {score},inference time is {e-s}")
    print("Finished!")

三、训练和测试结果

这里只给出部分结果，如下：
训练12epoch时，自己有事情请要用服务器，所以中断了训练，此时的准确率为：

用该模型测试3000张数据时，结果为：

注意：训练和测试代码中，设置的imgsize均为480，可以根据自己需求进行更改

---

总结

  以上就是本篇的全部内容，如有问题，欢迎评论区交流。