yolov5学习
一 基本使用
1 下载源代码
网址:https://github.com/ultralytics/yolov5
使用git clone到本地。
2 模型训练
使用pycharm打开yolov5文件夹,运行train.py脚本,会自动下载COCO128数据集,如果本地已有数据集,修改coco128.yaml里面数据集的路径。
也可以使用命令行进行训练,使用单gpu训练命令如下:
python train.py --batch 64 --data coco.yaml --weights yolov5s.pt --device 0
如果要从头开始训练,需要传入模型配置文件(yaml文件)
python train.py --data coco128.yaml --weights '' --cfg yolov5s.yaml --device 0
使用多gpu训练命令如下:https://github.com/ultralytics/yolov5/issues/475
python -m torch.distributed.run --nproc_per_node 2 train.py --batch 64 --data coco.yaml --weights yolov5s.pt --device 0,1
--nproc_per_node specifies how many GPUs you would like to use. In the example above, it is 2.
--batch is the total batch-size. It will be divided evenly to each GPU. In the example above, it is 64/2=32 per GPU.
3 模型测试
使用pycharm打开yolov5文件夹,运行detect.py脚本,模型会自动检测data/images路径下的所有图片和视频,
结果会保存至run文件夹下。根据需要修改模型权重路径,即weights参数,默认是将权重文件放到yolov5文
件夹。
预训练好的模型下载地址:https://github.com/ultralytics/yolov5/releases/tag/v6.1
也可以使用命令行进行训练,使用yolov5s预测单张图片命令行如下:
python detect.py --source img.jpg
使用yolov5s6模型预测单张图片命令行如下:
python detect.py --weights ./yolov5s6.pt --source ./zidane.jpg
使用yolov5s6模型预测文件夹中的所有图片命令行如下:
python detect.py --weights ./yolov5s6.pt --source ./data/images
4 模型导出
在终端窗口中运行export.py脚本,如果要在pycharm运行,需要配置参数。
将训练好的模型转成onnx,增加train选项,可以去掉后处理部分(detect layer)
python export.py --weights yolov5s6.pt --include onnx --train
5 平台使用
pt模型转成onnx模型后,首先使用onnx_post工具处理,命令行如下(需要在ubuntu虚拟机上运行):
./onnx_post ./yolov5s6.onnx ./new.onnx
之后使用编译器将onnx模型编译为bin文件,命令行如下(根据需要增删参数):
./compiler --output ./ --onnx ./new.onnx --image ./zidane.png --hybp --public_bin --mean '0.0,0.0,0.0' --norm '0.003921569,0.003921569,0.003921569'
二 使用自定义数据集进行训练
1 数据标注
使用labelimg工具进行数据标注,标注格式为yolo,安装使用待完善。
2 模型训练
在yolov5文件夹下新建datasets文件夹,然后在dataset文件夹下创建images文件夹和labels文件夹,之后分别在images文件夹和labels文件夹下创建train文件夹和val文件夹,目录结构如下图所示。
使用pycharm打开yolov5文件夹,
-
修改模型weights参数路径,默认在yolov5文件夹下使用yolov5s.pt进行训练。
-
修改cfg参数,使用的哪一个模型就使用哪一个模型的配置文件,配置文件在models文件夹下,yolov5x6保存在models/hub文件夹下
-
修改data参数,默认使用coco128.yaml配置文件,可以复制一份在原文件上进行修改,也可以在重命名一份进行修改(一般复制一份,直接在原文件上修改即可)。
-
修改coco128.yaml文件,将downlaod选项注释掉,根据需求修改剩下参数。
-
当显卡配置较好时,可以batch-size为32或者64或者128,当显卡配置一般,并且训练过程中提示显存不足时,可将batch-size设置为4
-
当存在多块显卡时,修改device参数,如存在两块显卡,device设置为“0,1”。
-
参数设置完成后,运行train脚本,开始训练即可。
3 模型测试
将测试图片保存至data/images文件夹下,运行detect脚本,需要修改weights参数为训练好的模型文件,如:
parser.add_argument(...default=ROOT / 'runs/train/exp/weights/best.pt'...)
根据需要修改source、data、imgsz等参数。
修改view-img、save-txt、save-conf、save-crop参数,可将测试结果保存成文件,保存内容为yolo格式,依次为x_center、y_center、w、h、confidence(坐标和宽高为归一化的值,如果要得到实际坐标,需要乘以原图片宽高)。
parser.add_argument('--view-img', default = True,action='store_true', help='show results')
parser.add_argument('--save-txt', default = True,action='store_true', help='save results to *.txt')
parser.add_argument('--save-conf', default = True,action='store_true', help='save confidences in --save-txt labels')
parser.add_argument('--save-crop', default = True,action='store_true', help='save cropped prediction boxes')
三 源码修改
1 pt转onnx模型
如果需要删除reshape和transpose层,将下面两行代码进行注释(yolov5-master\models\yolo.py)
def forward(self, x):
z = [] # inference output
for i in range(self.nl):
x[i] = self.m[i](x[i]) # conv
# 如果转onnx,需要将下面两行注释掉,删除reshape和transpose层,和编译器结果保持一致
bs, _, ny, nx = x[i].shape # x(bs,255,20,20) to x(bs,3,20,20,85)
x[i] = x[i].view(bs, self.na, self.no, ny, nx).permute(0, 1, 3, 4, 2).contiguous()
修改onnx模型输入shape。
yolov5输入为动态尺寸,不一定是正方形,需要根据输入模型的实际shape修改imgsz参数,之后导出模型。
def parse_opt():
parser = argparse.ArgumentParser()
parser.add_argument('--data', type=str, default=ROOT / 'data/coco128.yaml', help='dataset.yaml path')
parser.add_argument('--weights', nargs='+', type=str, default=ROOT / 'yolov5s.pt', help='model.pt path(s)')
parser.add_argument('--imgsz', '--img', '--img-size', nargs='+', type=int, default=[384,640], help='image (h, w)')
2 保存模型输入输出
保存输入模型的数据(yolov5-master\detect.py)
# Inference
visualize = increment_path(save_dir / Path(path).stem, mkdir=True) if visualize else False
# save input
testdata = im.detach().cpu().numpy()
import numpy as np
np.save('./yolov5_ZJ_input.npy', testdata)
testdata = testdata.reshape(-1)
np.savetxt('./yolov5_ZJ_input.txt', testdata, fmt="%s", delimiter='\n', newline='\n')
pred = model(im, augment=augment, visualize=visualize)
保存yolov5卷积后的输出(yolov5-master\models\yolo.py)
def forward(self, x):
z = [] # inference output
for i in range(self.nl):
x[i] = self.m[i](x[i]) # conv
# 保存yolov5卷积后的数据
arr = x[i].detach().cpu().numpy()
savepath = "output%d.npy" % (i)
np.save(savepath, arr)
arr = arr.reshape(-1)
savepath = "output%d.txt" % (i)
np.savetxt(savepath, arr, fmt="%s", delimiter='\n', newline='\n')
bs, _, ny, nx = x[i].shape # x(bs,255,20,20) to x(bs,3,20,20,85)
x[i] = x[i].view(bs, self.na, self.no, ny, nx).permute(0, 1, 3, 4, 2).contiguous()
3 修改resize插值方式
修改代码位置在yolov5-master\models\yolo.py
def _forward_once(self, x, profile=False, visualize=False):
y, dt = [], [] # outputs
for m in self.model:
if m.f != -1: # if not from previous layer
x = y[m.f] if isinstance(m.f, int) else [x if j == -1 else y[j] for j in m.f] # from earlier layers
if profile:
self._profile_one_layer(m, x, dt)
# change interpolation type
t = type(m)
if t is nn.Upsample:
m.mode = 'bilinear'
m.align_corners = True
x = m(x) # run
y.append(x if m.i in self.save else None) # save output
if visualize:
feature_visualization(x, m.type, m.i, save_dir=visualize)
return x
4 保存检测结果
需要注意的是保存的坐标是yolo格式,即归一化之后的坐标。
parser.add_argument('--save-txt', default = True,action='store_true', help='save results to *.txt')
parser.add_argument('--save-conf', default = True,action='store_true', help='save confidences in --save-txt labels')
5 打印检测结果
打印检测结果在原图上的坐标和置信度
# Write results
for *xyxy, conf, cls in reversed(det):
print("xyxy:",xyxy)
print("conf:",conf)
if save_txt: # Write to file
......
6 检测目标抠图保存
parser.add_argument('--save-crop',default = True,action='store_true', help='save cropped prediction boxes')
7 修改模型激活函数
将模型激活函数修改为线性激活函数(yolov5\models\common.py)
class Conv(nn.Module):
# Standard convolution
def __init__(self, c1, c2, k=1, s=1, p=None, g=1, act=True): # ch_in, ch_out, kernel, stride, padding, groups
super().__init__()
self.conv = nn.Conv2d(c1, c2, k, s, autopad(k, p), groups=g, bias=False)
self.bn = nn.BatchNorm2d(c2)
self.act = nn.SiLU() if act is True else (act if isinstance(act, nn.Module) else nn.Identity())
# 将激活函数改成线性激活函数
#self.act = nn.ReLU() if act is True else (act if isinstance(act, nn.Module) else nn.Identity())
class BottleneckCSP(nn.Module):
# CSP Bottleneck https://github.com/WongKinYiu/CrossStagePartialNetworks
def __init__(self, c1, c2, n=1, shortcut=True, g=1, e=0.5): # ch_in, ch_out, number, shortcut, groups, expansion
super().__init__()
c_ = int(c2 * e) # hidden channels
self.cv1 = Conv(c1, c_, 1, 1)
self.cv2 = nn.Conv2d(c1, c_, 1, 1, bias=False)
self.cv3 = nn.Conv2d(c_, c_, 1, 1, bias=False)
self.cv4 = Conv(2 * c_, c2, 1, 1)
self.bn = nn.BatchNorm2d(2 * c_) # applied to cat(cv2, cv3)
self.act = nn.SiLU()
#self.act = nn.ReLU()
self.m = nn.Sequential(*(Bottleneck(c_, c_, shortcut, g, e=1.0) for _ in range(n)))
8 查看模型参数量
修改img_size参数值,获取模型参数量,可以设置为不同的长宽
def info(self, verbose=False, img_size=1280): # print model information
model_info(self, verbose, img_size)
9 模型配置文件参数解释
backbone:
# [from, number, module, args]
# ch_in, ch_out, kernel, stride, padding, groups
[[-1, 1, Conv, [64, 6, 2, 2]], # 0-P1/2
[-1, 1, Conv, [128, 3, 2]], # 1-P2/4
[-1, 3, C3, [128]],
[-1, 1, Conv, [256, 3, 2]], # 3-P3/8
[-1, 6, C3, [256]],
[-1, 1, Conv, [512, 3, 2]], # 5-P4/16
[-1, 9, C3, [512]],
[-1, 1, Conv, [1024, 3, 2]], # 7-P5/32
[-1, 3, C3, [1024]],
[-1, 1, SPPF, [1024, 5]], # 9
]
10 断点续训
# Resume
if opt.resume and not check_wandb_resume(opt) and not opt.evolve: # resume an interrupted run
# 使用断点续训
# 如果resume是str,则表示传入的是模型路径
# 如果resume是true,则通过get_latest_run函数找到最近的pt文件
ckpt = opt.resume if isinstance(opt.resume, str) else get_latest_run() # specified or most recent path
assert os.path.isfile(ckpt), 'ERROR: --resume checkpoint does not exist'
with open(Path(ckpt).parent.parent / 'opt.yaml', errors='ignore') as f:
opt = argparse.Namespace(**yaml.safe_load(f)) # replace
opt.cfg, opt.weights, opt.resume = '', ckpt, True # reinstate
LOGGER.info(f'Resuming training from {ckpt}')
11 增加自定义激活函数
增加自定义激活函数(yolov5\utils\activations.py)
class FastSigmoid(nn.Module):
@staticmethod
def forward(x):
return x / (1 + abs(x))
class SiLU(nn.Module):
# SiLU activation https://arxiv.org/pdf/1606.08415.pdf
@staticmethod
def forward(x):
return x * torch.sigmoid(x)