赞
踩
前不久用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,可以根据自己需求进行更改
以上就是本篇的全部内容,如有问题,欢迎评论区交流。
Copyright © 2003-2013 www.wpsshop.cn 版权所有,并保留所有权利。