我的AI之路(45)--使用自己的数据集训练CenterNet

最新推荐文章于 2025-10-07 21:31:02 发布

原创

最新推荐文章于 2025-10-07 21:31:02 发布 · 3.8k 阅读

19 ·

CC 4.0 BY-SA版权

文章标签：

#深度学习 #CenterNet #pytorch0.4.1

本文详细介绍如何在Pytorch环境下配置并训练CenterNet模型，包括环境搭建、模型编译、数据集准备、训练过程及常见错误解决方法。

更新说明：

作者的源码: https://github.com/xingyizhou/CenterNet是基于pytorch0.4.1的(CUDA最高版本只能使用到CUDA9)，如果想使用pytorch1.0以上版本以支持使用CUDA10.0或以上版本，可以在下载作者的源码后，其他步骤照做，只是把DCNv2的源码用https://github.com/CharlesShang/DCNv2这里的源码替换掉就可以了，已在Pytorch1.2+RTX2080TI+CUDA10.0环境里训练了多个模型没任何问题：

cd CenterNet/src/lib/models/networks
rm -r DCNv2
git clone https://github.com/CharlesShang/DCNv2.git
cd DCNv2
python setup.py build develop

[原文]

CenterNet是anchor-free类型网络，具有识别精度高且速度快的特点，根据作者的论文中列出的数据来看，指标综合考虑来看比较牛了：

最后那个CenterNet-HG，也就是backbone使用的Hourglass-104网络的AP值只比FSAF低一点了(但是FSAF目前貌似还没有源码放出来),比YOLO序列和RCNN序列都强很多，虽然FPS自有7.8，但是对一般实时性要求不是很高的视频检测也够用了，所以拿来试试。

首先下载作者的源码: git clone https://github.com/xingyizhou/CenterNet.git，根据安装说明:https://github.com/xingyizhou/CenterNet/blob/master/readme/INSTALL.md，环境和工具软件是:

Ubuntu 16.04, with Anaconda Python 3.6 and PyTorch v0.4.1

他这个源码是使用的pytorch0.41版写的，由于pytorch0.41支持的CUDA最高版本是CUDA9，不支持我们的服务器上目前安装的CUDA10.0或CUDA10.1,我先是试了一下使用conda创建隔离环境后安装支持CUDA10的pytorch1.3或pytorch1.0.0，然后跑了一下，结果报错，说是CenterNet中有API是不支持的了(后面再说)，但在公共服务器上又不好随便乱装CUDA(安装过CUDA的应该知道它的厉害，很能折腾人，装得不对服务器登录进不去、黑屏之类的问题让人三思)，于是想到还是使用docker最好，首先到hub.docker.com上拉取个pytorch0.4.1+CUDA9.0的devel版镜像:

docker pull pytorch/pytorch:0.4.1-cuda9-cudnn7-devel

然后运行创建实例(进入容器内部后默认的初始路径是/workspace，所以把下载了CenterNet源码的目录work_pytorch映射到/workspace，并预留端口12000的映射，以备后面有需要时对模型做server端封装调用，并带上ipc=host参数，以防止做多GPU分布式训练的过程中出现共享内存不足的错误):

nvidia-docker run --ipc=host -d -it --name pytorch0.41 -v /home/fychen/AI/work_pytorch:/workspace -p 12000:12000 pytorch/pytorch:0.4.1-cuda9-cudnn7-devel bash

进入容器后，执行下面的修改(容器内的pytorch安装在/opt/conda路径下)把torch.nn.functional.py里1254行的torch.backends.cuddn.enabled改为False：

sed -i "1254s/torch\.backends\.cudnn\.enabled/False/g" /opt/conda/lib/python3.6/site-packages/torch/nn/functional.py

然后，依次执行下面的命令安装pycocotools:

git clone https://github.com/cocodataset/cocoapi.git
cd cocoapi/PythonAPI
pip install cython
make
python setup.py install --user

再依次执行下面的命令完成CenterNet下面的部分代码的编译:

cd /workspace/CenterNet
pip install -r requirements.txt

cd src/lib/models/networks/DCNv2
./make.sh

cd /workspace/CenterNet/src/lib/external
make

再安装一些跑CenterNet需要的支持包(不安装这些包会报错)：

apt-get install libglib2.0-dev libsm6 libxrender1 libxext6

然后下载对应的预训练模型，我要使用的backbone是hour-glass，模型训练后用来做物体检测，根据https://github.com/xingyizhou/CenterNet/blob/master/readme/MODEL_ZOO.md :