Detectron2: register META_ARCH

[yimengli46]

I implemented my own semantic segmentation model and registered it as META_ARCH.
And I use it in the config file.

@META_ARCH_REGISTRY.register()
class BayesianSemanticSegmentor(nn.Module):

However, when I try to train this model through this config file, I get this key error.
KeyError: "No object named 'BayesianSemanticSegmentor' found in 'META_ARCH' registry!"
I print out the registered keys in META_ARCH_REGISTRY, my model is not there.
META_ARCH_REGISTRY.keys = dict_keys(['SemanticSegmentor', 'PanopticFPN', 'GeneralizedRCNN', 'ProposalNetwork', 'RetinaNet'])

So what step am I missing here?

[ppwwyyxx]

My guess is the code is not executed (e.g imported).
However we'll need details about the problem following the issue template if you need any further help.

[yimengli46]

Hi Yuxin, I have the following files.
_semantic.yaml_

_BASE_: "Base-RCNN-FPN.yaml"
MODEL:
  META_ARCHITECTURE: "BayesianSemanticSegmentor"
  WEIGHTS: "detectron2://ImageNetPretrained/MSRA/R-50.pkl"
  RESNETS:
    DEPTH: 50
DATASETS:
  TRAIN: ("coco_2017_train_panoptic_stuffonly",)
  TEST: ("coco_2017_val_panoptic_stuffonly",)
INPUT:
  MIN_SIZE_TRAIN: (640, 672, 704, 736, 768, 800)

_init.py_

from .bayesian_semantic_seg import (
    BayesianSemanticSegmentor, 
    build_sem_seg_head, 
    SemSegFPNDropoutHead,
    )

_bayesian_semantic_seg.py_

from detectron2.layers import Conv2d, ShapeSpec
from detectron2.structures import ImageList
from detectron2.utils.registry import Registry

from detectron2.modeling import build_backbone, META_ARCH_REGISTRY
from detectron2.modeling.postprocessing import sem_seg_postprocess

from .conv2d_dropout import Conv2d_Dropout

__all__ = ["BayesianSemanticSegmentor", "SEM_SEG_HEADS_REGISTRY", "build_sem_seg_head", "SemSegFPNDropoutHead"]

SEM_SEG_HEADS_REGISTRY = Registry("SEM_SEG_HEADS")
"""
Registry for semantic segmentation heads, which make semantic segmentation predictions
from feature maps.
"""

@META_ARCH_REGISTRY.register()
class BayesianSemanticSegmentor(nn.Module):
    """
    Main class for semantic segmentation architectures.
    """

    def __init__(self, cfg):
        super().__init__()

        self.device = torch.device(cfg.MODEL.DEVICE)

        self.backbone = build_backbone(cfg)
        self.sem_seg_head = build_sem_seg_head(cfg, self.backbone.output_shape())

        pixel_mean = torch.Tensor(cfg.MODEL.PIXEL_MEAN).to(self.device).view(-1, 1, 1)
        pixel_std = torch.Tensor(cfg.MODEL.PIXEL_STD).to(self.device).view(-1, 1, 1)
        self.normalizer = lambda x: (x - pixel_mean) / pixel_std

        self.to(self.device)

_train_sseg.py_

epoch = 1
thresh = 0.7
cfg = get_cfg()
cfg.merge_from_file('configs/semantic_R_50_FPN.yaml')
cfg.DATASETS.TRAIN = ('mapillary_sseg_train', ) # the comma is necessary
cfg.DATASETS.TEST = ()
cfg.DATALOADER.NUM_WORKERS = 8
cfg.MODEL.WEIGHTS = 'model_weights/coco_panopticSeg_R_50_FPN_1x.pkl'
cfg.MODEL.SEM_SEG_HEAD.NAME = 'SemSegFPNDropoutHead'
cfg.MODEL.SEM_SEG_HEAD.IGNORE_VALUE = 255
cfg.MODEL.SEM_SEG_HEAD.NUM_CLASSES = 66
cfg.SOLVER.IMS_PER_BATCH = 8
cfg.SOLVER.BASE_LR = 0.00025
cfg.SOLVER.MAX_ITER = 2200 * epoch
cfg.OUTPUT_DIR = 'trained_model/bayesianSeg_Mapillary_model_epoch_{}'.format(epoch)

os.makedirs(cfg.OUTPUT_DIR, exist_ok=True)
trainer = DefaultTrainer(cfg)
trainer.resume_or_load(resume=False)
trainer.train()

When I run _train_sseg.py_, the error is
KeyError: "No object named 'BayesianSemanticSegmentor' found in 'META_ARCH' registry!"

I tried to compare the code with the _TensorMask_ project, as _TensorMask_ also built a new _META_ARCH_, but I couldn't find what code is missing here.

[ppwwyyxx]

As I said above, you need to import the file bayesian_semantic_seg.

[yimengli46]

Thanks. Problem solved. Even though I don't fully understand why I have to explicitly import _bayesian_semantic_seg_ in _train_sseg.py_, as I don't call any functions directly in _bayesian_semantic_seg_.