在detectron2/data/datasets/builtin.py中可以看到在DatasetCatelog上各个数据集的注册,其中_root即为数据集的基地址,代码指明,_root要么是DETECTRON2_DATASETS,要么是datasets/,因此想要修改数据集位置,可以通过在命令行运行export DETECTRON2_DATASETS=/xx/xx/xx来修改。
下面包括基于VOC的Faster RCNN的训练和测试:代码基于train_net.py完成
from detectron2.utils.logger import setup_logger
setup_logger()
from detectron2 import model_zoo
from detectron2.config import get_cfg
from tools.train_net import Trainer
from detectron2.checkpoint import DetectionCheckpointer
import os
os.system("export DETECTRON2_DATASETS=/media/zhou/data/VOC/VOCdevkit/")
os.system("echo $DETECTRON2_DATASETS")
def train():
cfg.DATASETS.NUM_WORKERS = 1
cfg.SOLVER.IMS_PER_BATCH = 4
cfg.SOLVER.BASE_LR = 0.01
cfg.MODEL.ROI_HEADS.NUM_CLASSES = 20
# os.makedirs(cfg.OUTPUT_DIR, exist_ok=True)
trainer = Trainer(cfg)
trainer.resume_or_load(resume=False)
trainer.train()
def test():
cfg.MODEL.WEIGHTS = './output/model_final.pth'
model = Trainer.build_model(cfg)
DetectionCheckpointer(model, save_dir=cfg.OUTPUT_DIR).resume_or_load(
cfg.MODEL.WEIGHTS, resume=False
)
Trainer.test(cfg, model)
if __name__ == '__main__':
cfg = get_cfg()
cfg.merge_from_file(model_zoo.get_config_file("../configs/PascalVOC-Detection/faster_rcnn_R_50_FPN.yaml"))
官方没有跑R50-FPN,VOC,因此只能和R50-C4,VOC作比较。结果相差两个点,可以接受,一个是因为模型框架不同,另一个是因为我使用的batch_size = 4,偏小。
重新训练了Faster rcnn c4,得到的结果较差,可能是因为没有使用官方的训练参数,以及batch_size = 2,较小的原因