文章目录
- 前言
- 一、数据摆放
- 二、训练
- 二、测试
- 三、训练和测试结果
- 总结
前言
前不久用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，可以根据自己需求进行更改
总结
以上就是本篇的全部内容，如有问题，欢迎评论区交流。
