当前位置: 首页 > >

maskrcnn-benchmark 训练自己的数据集

发布时间:

代码修改
1、使用预训练模型需要修改部分代码,主要由于自己的数据集类别和预训练的81类对不上。
修改maskrcnn-benchmask/config/default.py:
_C.MODEL.ROI_BOX_HEAD.NUM_CLASSES修改为自己的数据集类别加1
_C.SOLVER.CHECKPOINT_PERIOD修改为10000(可选),设置大点是为了减少保存的模型数
修改maskrcnn-benchmask/modeling/roi_heads/box_head/roi_box_predictors.py:
由于我们最终输出的类别和预训练模型不一致,所以我们需要对最后的几层网络修改一下名字,以免读取模型错误。




2、修改maskrcnn-benchmask/utils/checkpoint.py,在其中注释读取optimizer和scheduler的代码,这也是由于数据集类别对不上导致的。




3、修改maskrcnn-benchmask/data/datasets/voc.py,修改其中的CLASSES为你自己的数据集的类别。



友情链接: