CO-DETR利用coco数据集训练和推理过程,参考链接
Co-DETR训练自己的数据集
文章目录
- 前言
- 训练过程
- 推理过程
- 总结
前言
环境:PyTorch 1.11.0 Python 3.8(ubuntu20.04) Cuda 11.3
先是在github上下载CO-DETR模型
!git clone https://github.com/Sense-X/Co-DETR.git
%cd Co-DETR
然后加载所需库
!pip install -r requirements.txt
安装mmcv等(注意mmcv应该是1.6.1版本及以上)
!pip install -U openmim
!mim install mmcv-full==1.6.1
!pip install timm==0.6.11 mmdet==2.25.3
因为出现了mmdetection 报错 TypeError: FormatCode() got an unexpected keyword argument ‘verify‘问题,用一下方案解决: yapf版本过高,目前版本为 0.40.2,重装低版本yapf即可
!pip uninstall yapf
!pip install yapf==0.40.1
然后把解压好的coco数据集放到/Co-DETR/data/coco/
目录下,如下图
如果出现ModuleNotFoundError: No module named ‘projects‘
错误,在相应python文件中(一般是train.py)添加
import sys
sys.path.append('你的项目的绝对路径')
例如:
#/Co-DETR/tools/train.py里面修改
import sys
sys.path.append('/root/autodl-tmp/Co-DETR')
from projects import *
训练过程
在/Co-DETR
目录下,终端输入:
bash tools/dist_train.sh projects/configs/co_deformable_detr/co_deformable_detr_r50_1x_coco.py 2 /root/autodl-tmp/Co-DETR
因为是分布式训练,需要用到两张以上显卡,比如两张4090(第三个参数是gpu数量),第一个参数是利用Pytorch的torch.distributed 实现单机多卡分布式训练的shell脚本文件,第二个是模型配置文件的位置,采用的是co_deformable_detr_r50_1x_coco.py
模型,第四个参数是跑出的权重放置的位置。
运行的日志如下图所示:
下图是正常运行时终端正在运行的情况,正在进行第一轮训练:
然后等代码跑完12轮就行了,两张4090三四个小时跑完一轮,如下图
应该可以通过修改如下红色区域的值来修改跑的轮次
推理过程
可以用上面跑的权重,或者想简单点直接在官方代码中找到对应模型的权重,下载到服务器(gpu)里
用scp或者wget等方式下载一张行车记录仪的记录图片到/Co-DETR/deno
目录下,比如test.png,运行下面命令,在当前目录下会出现识别后的out2.png图像。
!python demo/image_demo.py demo/test.png \
projects/configs/co_deformable_detr/co_deformable_detr_r50_1x_coco.py \
checkpoints/co_deformable_detr_r50_1x_coco.pth \
--device cuda \
--out-file out2.png
第一个参数是图片位置,第二个参数是模型配置位置,第三个参数是权重位置,第五个参数收识别图像的输出位置。
该图像是运行官网的权重后的结果
下图是运行上面第一轮训练后的latest.pth权重的结果
明显才训练一轮的识别结果更差,并且还有远处的一辆truck没识别到
总结
不记得运行过程有没有遇到其他问题了,但是基本上都能搜出解决方案,之前一直有问题感觉主要是没用多张卡来跑