CenterNet 训练自己的数据集

github地址:https://github.com/Duankaiwen/CenterNet

论文:https://arxiv.org/abs/1904.08189

1、在github上下载,配置好环境。

2、准备数据

把自己的数据转化为coco的格式,网上有很多工具可以下载使用。我的数据是yolo格式的,需要的话可以提供我的yolo to coco的代码。yolo to coco数据转化小工具:https://github.com/surserrr/yolo_to_coco (随手star啦谢谢~~

数据分成train和val。 图片文件夹名字改为trainval2014和minival2014,放到CenterNet-master/data/coco/images中; json文件名为instances_trainval2014.json和instances_minival2014.json,放到CenterNet-master/data/coco/annotations中。

注意:如果你在训练之前,用coco数据集测试了模型,那么把CenterNet-master/cache/coco_minival2014.pkl删掉!如果没有训练过就忽略。(因为你在第一次运行的时候,代码会把coco数据的instances转化为它要用的格式,下一次用的时候就会直接读取。如果你没删掉,但你训练自己的数据集的时候,模型在val的时候会自动读取已经存在的coco的数据)

3、修改参数。

我数据集的类别只有1类,GPU1个。

必须修改的参数:

    1)models/CenterNet-52.py或者models/CenterNet-104.py(看你要训练哪一个模型,下面省略),中第132行 out_dim =80  ==>  out_dim =1 (类别数)

    2)config/CenterNet-52.py第4行 "batch_size":48 ==>  "batch_size": 6 (我的显卡是Titan X ,按自己的显存设置,可以设置了试一下,运行的时候如果溢出了,再改小。)

    3)  config/CenterNet-52.py第22行 "chunk_size": [6,6,6,6,6,6,6,6] ==> "chunk_size": [6] 。chunk_size长度是gpu数量,chunk_size的sum是batchsize。原文的环境是8块显卡,batchsize48,每块显卡分配batchsize 6 。 我只有一块显卡,所以设置为[6]。 如果你有两块显卡,batchsize是12,那么chunk_size为[6,6];如果你一块显卡,batchsize是2,那么chunk_size为[2].

    4)  config/CenterNet-52.py第45行 “categories":80  ==>   “categories":1 (类别数)

    5)db/detection.py 第8行self._configs[ “categories"] =80 ==> self._configs[ “categories"] = 1 (类别数)

    6)db/coco.py 第47行 self._cat_ids 改为 self._cat_ids = [1] (看你有几个类别)

    7)db/detection.py 第72行 num_classes=81 ==> num_classes=2 (类别数+1)

选择修改的参数:

config/CenterNet-52.py中 18行 -20行

"max_iter": 40000,   #迭代次数
"stepsize": 10000,   #learning rate下降次数
"snapshot": 1000,    #模型每隔snapshot保存一次

4、训练

python train.py CenterNet-52

就可以训练啦,30小时以后看看训练效果怎么样。

5、更新

失败了。。。上面的训练是没有问题的。

失败的原因是和我的数据集不合适。

CenterNet的核心就是预测角点,再预测中心点,判断中心点和角点是不是一个类型来确定这个检测框对不对。。。

但是我的数据集是一个在刚体表面检测划痕的,,划痕和普通检测还是又不一样,可能这个划痕根本没有穿过中心点,所以上述的约束条件就很难了。。失败了。不过还是从这个论文里学到了。

之前有了解过一段时间的人体姿态检测,一直觉得bottom-up的方法比up-down的方法要高超,现在检测出的一系列基于keypoint的,其实和姿态检测里bottom-up差不太多。。如果有条件的话真想好好研究一下。

你可能感兴趣的:(计算机视觉)