cfg.DATASETS.TEST = () cfg.DATALOADER.NUM_WORKERS = num_workers cfg.MODEL.WEIGHTS = model_zoo.get_checkpoint_url("COCO-InstanceSegmentation/mask_rcnn_R_50_FPN_3x.yaml") cfg.SOLVER.IMS_PER_BATCH = ims_per_batch cfg.SOLVER.BASE_LR = model_lr cfg.SOLVER.MAX_ITER = max_train_iter cfg.MODEL.ROI_HEADS.BATCH_SIZE_PER_IMAGE = bach_size_per_img cfg.MODEL.ROI_HEADS.NUM_CLASSES = num_labels cfg.OUTPUT_DIR = "{}/maskrcnn_{}_{}_{}_{}_{}".format(work_root,num_labels,max_train_iter,bach_size_per_img,ims_per_batch,model_lr) os.makedirs(cfg.OUTPUT_DIR, exist_ok=True) ####### 训练 # trainer = CustomTrainer(cfg) trainer = DefaultTrainer(cfg) train_data_loader = trainer.build_train_loader(cfg) ######## 可视化 # data_iter = iter(train_data_loader) # batch = next(data_iter) # rows, cols = 2, 2 # plt.figure(figsize=(20,20)) # # for i, per_image in enumerate(batch[:4]): # # plt.subplot(rows, cols, i+1) # # # Pytorch tensor is in (C, H, W) format # img = per_image["image"].permute(1, 2, 0).cpu().detach().numpy() # img = utils.convert_image_to_rgb(img, cfg.INPUT.FORMAT) #
cfg.merge_from_file(model_zoo.get_config_file("COCO-InstanceSegmentation/mask_rcnn_X_101_32x8d_FPN_3x.yaml")) cfg.DATALOADER.NUM_WORKERS = 2 cfg.MODEL.WEIGHTS = model_zoo.get_checkpoint_url("COCO-InstanceSegmentation/mask_rcnn_X_101_32x8d_FPN_3x.yaml") # Let training initialize from model zoo cfg.SOLVER.IMS_PER_BATCH = 2 cfg.SOLVER.BASE_LR = 0.00025 # pick a good LR cfg.SOLVER.MAX_ITER = 1000 # 300 iterations seems good enough for this toy dataset; you may need to train longer for a practical dataset cfg.MODEL.ROI_HEADS.BATCH_SIZE_PER_IMAGE = 256 # faster, and good enough for this toy dataset (default: 512) cfg.MODEL.ROI_HEADS.NUM_CLASSES = 1 # only has one class (ballon) cfg.DATASETS.TRAIN = ('my_dataset',) cfg.DATASETS.TEST = () os.makedirs(cfg.OUTPUT_DIR, exist_ok=True) trainer = DefaultTrainer(cfg) trainer.build_train_loader = build_detection_train_loader(cfg, mapper=train_mapper) trainer.resume_or_load(resume=False) trainer.train() from contextlib import redirect_stdout #ouput config with open('test.yaml', 'w') as f: with redirect_stdout(f): print(cfg.dump()) from detectron2.evaluation import COCOEvaluator, inference_on_dataset from detectron2.data import build_detection_test_loader evaluator = COCOEvaluator("test_dataset", cfg, False, output_dir="./output/") val_loader = build_detection_test_loader(cfg, "test_dataset") inference_on_dataset(trainer.model, val_loader, evaluator)