文章目录
- 前言
- 1、准备工作
- 2、训练模型
-
- 2.1.可能报的错:AttributeError: Cant get attribute SPPF on module models.common
- 2.2.训练模型
- 2.3.测试模型
- 可视化
- 总结
前言
扔掉学术偏见,拥抱工程化的Yolo。由于一直没有使用过yolo,因此本文旨在从0基础配置yolo环境并训练和测试。
论文地址
代码地址
1、准备工作
- 首先下载coco2017数据集,怎么下载这个我就不详细介绍了,但注意的是,虽然原始coco数据集的文件夹名称也为coco,但为了和yolo所使用的进行区分,请将coco数据集命名为coco2017。
- 然后下载原始coco格式的标签并解压,解压后请务必将person_keypoints_val2017.json文件拷贝到coco/annotations/下,否则后续评估报错!
- 新建一个coco_kpts文件夹,下载yolo格式的关键点检测的标签并解压到coco_kpts文件夹内(谷歌云盘,没有梯子的自己想想办法把,别私聊问我要了);同时创建软连接或者将coco2017中的images和annotations两个文件夹拷贝到coco_kpts下面。最终coco_kpts/如下图:先忽略train2017.cache,后续在训练时会生成。
最后下载代码并解压至和coco_kpts同级目录下即可。
在准备完成后,最终的一个文件树为:
至此,便可以训练yolo-pose。
2、训练模型
2.1.可能报的错:AttributeError: Cant get attribute SPPF on module models.common
参考链接
在model/common.py文件中加上如下代码即可。
import warnings
class SPPF(nn.Module):
# Spatial Pyramid Pooling - Fast (SPPF) layer for YOLOv5 by Glenn Jocher
def __init__(self, c1, c2, k=5): # equivalent to SPP(k=(5, 9, 13))
super().__init__()
c_ = c1 // 2 # hidden channels
self.cv1 = Conv(c1, c_, 1, 1)
self.cv2 = Conv(c_ * 4, c2, 1, 1)
self.m = nn.MaxPool2d(kernel_size=k, stride=1, padding=k // 2)
def forward(self, x):
x = self.cv1(x)
with warnings.catch_warnings():
warnings.simplefilter('ignore') # suppress torch 1.9.0 max_pool2d() warning
y1 = self.m(x)
y2 = self.m(y1)
return self.cv2(torch.cat([x, y1, y2, self.m(y2)], 1))
2.2.训练模型
python train.py --data coco_kpts.yaml --cfg yolov5s6_kpts.yaml --batch-size 64 --img 640 --kpt-label
需要训练300个epoch,等训练完我会补充下结果,目前来看,效果果然好,3个epoch情况下达到了oks达到了38.3,恐怖。:
训练了100epoch,oks在测试集精度为46。官方汇报的是57,训了300epoch。
2.3.测试模型
首先在readme中下载一个训练好的权重,如下图,点击那个Yolo5s6_pose_640,将下载好的.pt放到edgeai-yolov5-yolo-pose文件夹内。
测试命令:
python test.py --data coco_kpts.yaml --img 640 --conf 0.001 --iou 0.65 --weights "/path/to/edgeai-yolov5-yolo-pose/last.pt" --kpt-label
测试结果:59.9比官方汇报的高两个点。
可视化
测试/训练的可视化图都在runs文件夹内,我这里放几张测试图例:
总结
有问题欢迎+vx: wulele2541612007,后续有空会出yolo-pose的源码解读,敬请期待。
近期评论