How to setup NN structure that have shared weight

[songanz]

Hi, I am working on a project that will need to have multiple input images to the NN and do the detection on them. My problem is how can I first extract features of these images through ONE backbone and input them to other module all together in Detectron2? The size of the whole NN is critical in my application. So I don't want to "copy" that backbone several times.

Is is possible to add a wrapper in Detectron2 like this? If so, how should I register this kind of wrapper and setup the config file in Detectron2?

class MyModel(nn.Module):
    def __init__(self):
        self.base = ...
        self.head_A = ...
        self.head_B = ...

    def forward(self, input1, input2):
        return self.head_A(self.base(input1)), self.head_B(self.base(input2))