Support of Detectron2 models about nebuly HOT 1 CLOSED

class BaseModelWrapper(torch.nn.Module):
    def __init__(self, core_model, output_dict):
        super().__init__()
        self.core_model = core_model
        self.output_names = [key for key in output_dict.keys()]
    
    def forward(self, *args, **kwargs):
        res = self.core_model(*args, **kwargs)
        return tuple(res[key] for key in self.output_names)


class OptimizedWrapper(torch.nn.Module):
    def __init__(self, optimized_model, output_keys):
        super().__init__()
        self.optimized_model = optimized_model
        self.output_keys = output_keys
    
    def forward(self, *args):
        res = self.optimized_model(*args)
        return {key: value for key, value in zip(self.output_keys, res)}

import copy
from detectron2 import model_zoo 
from nebullvm import optimize_torch_model

config_path = "COCO-InstanceSegmentation/mask_rcnn_R_50_FPN_3x.yaml"
model = model_zoo.get(config_path, trained=True) 
model.eval()
model_backbone = copy.deepcopy(model.backbone)
res = model_backbone(torch.randn(1, 3, 256, 256))  # needed for getting the output_keys
backbone_wrapper = BaseModelWrapper(model_backbone, res)
optimized_model = optimize_torch_model(backbone_wrapper, batch_size=1, input_sizes=[(3, 256, 256)], save_dir="./")
optimized_backbone = OptimizedWrapper(optimized_model, backbone_wrapper.output_names)
# finally replace the old backbone with the optimised one
model.backbone = optimized_backbone