cfg parameters don't overwrite custom ROIHeads implementation

[eugeniotonanzi]

Hi,
I'm working with Detectron2 Faster R-CNN models. What I need to do is to extract features that are 2048-dimensional from Regions of Interest, rather than the default 1024-dimensional features. So, starting from a model from the model zoo, I registered a new custom ROIHeads class which should be the same as StandardROIHeads class but replacing the linear layers in order to have dimension 2048:

from detectron2.config import get_cfg
from detectron2 import model_zoo

from detectron2.modeling import ShapeSpec
from detectron2.modeling.roi_heads.box_head import FastRCNNConvFCHead
from detectron2.modeling.roi_heads.fast_rcnn import FastRCNNOutputLayers
from detectron2.modeling.roi_heads import StandardROIHeads
from detectron2.modeling import ROI_HEADS_REGISTRY

cfg = get_cfg()
cfg.merge_from_file(model_zoo.get_config_file("COCO-Detection/faster_rcnn_R_101_FPN_3x.yaml"))

box_head_2048 = FastRCNNConvFCHead(cfg, input_shape=ShapeSpec(channels=256, height=7, width=7), fc_dims=[2048, 2048])

box_predictor_2048 = FastRCNNOutputLayers(cfg, input_shape=ShapeSpec(channels=2048))

@ROI_HEADS_REGISTRY.register()
class StandardROIHeads2048(StandardROIHeads):
def init(self, cfg, input_shape):
super().init(cfg, input_shape, box_head=box_head_2048, box_predictor=box_predictor_2048)

cfg.MODEL.ROI_HEADS.NAME = 'MyStandardROIHeads2048'

The problem is that if I try to overwrite its parameters using cfg, nothing happens. That is, if I run for example

cfg.MODEL.ROI_HEADS.NUM_CLASSES = 1
model = build_model(cfg)

to train the model on the balloon dataset from the Detectron2 tutorial, the model ROIHeads has still a final output with 81 classes, as the default. Instead, if I use the standard model without my custom ROIHeads, using the previous command works and overwrites the cfg parameters, making the output 2-dimensional rather than 81.

I don't get where I'm doing it wrong, since my custom implementation just subclasses the StandardROIHeads used by the default model.