源码下载:https://github.com/ultralytics/ultralytics
权重下载:https://github.com/ultralytics/assets/releases
为了更好的理解yolov8代码细节,复现时没有采用ultralytics预定义好的指令集,而是在根目录下创建main.py作为源码起点,从而可以更好的进行断点调试理解内部数据流,具体实现如下:
###################################################################
# flag = 0, predict detect
# flag = 1, predict pose
# flag = 2, train detect
# flag = 3, train pose
# flag = 10, transform .pt file to .onnx file
# flag = 11, add shape into onnx show by netron
###################################################################
flag = 3
if flag==0:
from ultralytics.utils import ASSETS
from ultralytics.models.yolo.detect import DetectionPredictor
args = dict(model='yolov8s.pt', source=ASSETS)
predictor = DetectionPredictor(overrides=args)
predictor.predict_cli()
elif flag==1:
from ultralytics.utils import ASSETS
from ultralytics.models.yolo.pose import PosePredictor
args = dict(model='yolov8s-pose.pt', source=ASSETS)
predictor = PosePredictor(overrides=args)
predictor.predict_cli()
elif flag==2:
from ultralytics.models.yolo.detect import DetectionTrainer
args = dict(model='yolov8s.pt', data='coco8.yaml', epochs=3)
trainer = DetectionTrainer(overrides=args)
trainer.train()
elif flag==3:
from ultralytics.models.yolo.pose import PoseTrainer
args = dict(model='yolov8s-pose.pt', data='coco8-pose.yaml', epochs=3)
trainer = PoseTrainer(overrides=args)
trainer.train()
elif flag==10:
from ultralytics import YOLO
model = YOLO('yolov8s-pose.pt')
results = model.export(format='onnx')
elif flag==11:
import onnx
from onnx import shape_inference
model_onnx = r'/data/deeplearningalgo/ultralytics/yolov8s-pose.onnx'
onnx.save(onnx.shape_inference.infer_shapes(onnx.load(model_onnx)), model_onnx)
通过设置flag值执行yolov8的不同功能操作,其中(0,1,2,3)分别代表检测任务预测、关键点任务预测、检测任务训练、关键点任务训练,flag==10时输入为预下载好的权重文件或者本地训练保存好的权重文件,输出为转换好的onnx文件,如果更加形象的了解模型结构,可以使用netron工具打开转换好的onnx文件,直接浏览器打开Netron,打开对应的模型文件,仅通过flag==10进行转换生成的onnx文件打开无法显示输出shape,对网络结构的学习影响不小,因此还需要在此基础上继续执行flag==11的转换,才可生成带shape的模型结构图,由于学习过程中找了很久都没找到合适的网络图,所以我在这里把生成的onnx文件也上传到了资源,大家可以直接下载打开即可学习。
yolov8工程整体代码封装的更好,和之前看yolox的代码时感受不同(可以通过vscode一路跳转),由于yolov8功能更加丰富(包括检测、分割、关键点等),通过读取不同的模型yaml文件来(路径:/your_code_path/ultralytics/cfg/v8/)配置模型参数,其中yolov8-pose.yaml如下:
# Parameters
nc: 1 # number of classes
kpt_shape: [17, 3] # number of keypoints, number of dims (2 for x,y or 3 for x,y,visible)
scales: # model compound scaling constants, i.e. 'model=yolov8n-pose.yaml' will call yolov8-pose.yaml with scale 'n'
# [depth, width, max_channels]
n: [0.33, 0.25, 1024]
s: [0.33, 0.50, 1024]
m: [0.67, 0.75, 768]
l: [1.00, 1.00, 512]
x: [1.00, 1.25, 512]
# YOLOv8.0n backbone
backbone:
# [from, repeats, module, args]
- [-1, 1, Conv, [64, 3, 2]] # 0-P1/2
- [-1, 1, Conv, [128, 3, 2]] # 1-P2/4
- [-1, 3, C2f, [128, True]]
- [-1, 1, Conv, [256, 3, 2]] # 3-P3/8
- [-1, 6, C2f, [256, True]]
- [-1, 1, Conv, [512, 3, 2]] # 5-P4/16
- [-1, 6, C2f, [512, True]]
- [-1, 1, Conv, [1024, 3, 2]] # 7-P5/32
- [-1, 3, C2f, [1024, True]]
- [-1, 1, SPPF, [1024, 5]] # 9
# YOLOv8.0n head
head:
- [-1, 1, nn.Upsample, [None, 2, 'nearest']]
- [[-1, 6], 1, Concat, [1]] # cat backbone P4
- [-1, 3, C2f, [512]] # 12
- [-1, 1, nn.Upsample, [None, 2, 'nearest']]
- [[-1, 4], 1, Concat, [1]] # cat backbone P3
- [-1, 3, C2f, [256]] # 15 (P3/8-small)
- [-1, 1, Conv, [256, 3, 2]]
- [[-1, 12], 1, Concat, [1]] # cat head P4
- [-1, 3, C2f, [512]] # 18 (P4/16-medium)
- [-1, 1, Conv, [512, 3, 2]]
- [[-1, 9], 1, Concat, [1]] # cat head P5
- [-1, 3, C2f, [1024]] # 21 (P5/32-large)
- [[15, 18, 21], 1, Pose, [nc, kpt_shape]] # Pose(P3, P4, P5)
由于yolov8工程封装性比较好,也就导致代码阅读起来会比较困难,跳着跳着就不知道下一步要执行什么步骤了,又或者忘了之前从哪跳转过来的,基于此,为了让大家少走弯路,我将模型的训练过程绘制成了UML时序图,大家照着图去一步步看yolov8源码一定会事半功倍
刚入门没多久,一点学习心得,希望能够帮助到大家,如果有帮助麻烦帮忙点击下收藏关注,也欢迎大家批评指正交流~