How to save Detectron2 model as a vanilla pytorch model?

Question

I have a Faster-RCNN model trained with Detectron2. Model weights are saved as model.pth.

I have my pickled config.yml file and there are a couple of ways to load this model:

from detectron2.modeling import build_model
from detectron2.checkpoint import DetectionCheckpointer

cfg = get_cfg()
config_name = "config.yml" 
cfg.merge_from_file(config_name)

cfg.MODEL.WEIGHTS = './model.pth'
model = DefaultPredictor(cfg)

OR

model_ = build_model(cfg) 
model = DetectionCheckpointer(model_).load("./model.pth")

Also, you can get predictions from this model individually as given in official documentation:

image = np.array(Image.open('page4.jpg'))[:,:,::-1] # RGB to BGR format
tensor_image = torch.from_numpy(image.copy()).permute(2, 0, 1) # B, channels, W, H


with torch.no_grad():
    output = torch_model([{"image":tensor_image}])

running the following commands:

print(type(model))
print(type(model.model))
print(type(model.model.backbone))

Gives you:

<class 'detectron2.engine.defaults.DefaultPredictor'>
<class 'detectron2.modeling.meta_arch.rcnn.GeneralizedRCNN'>
<class 'detectron2.modeling.backbone.fpn.FPN'>

Problem: I want to use GradCam for model explainability and it uses pytorch models as given in this tutorial

How can I turn detectron2 model in vanilla pytorch model?

I have tried:

torch.save(model.model.state_dict(), "torch_weights.pth")
torch.save(model.model, "torch_model.pth")


from torchvision.models.detection import fasterrcnn_resnet50_fpn

dummy = fasterrcnn_resnet50_fpn(pretrained=False, num_classes=1)
# dummy.load_state_dict(torch.load('./model.pth', map_location = 'cpu')) 
dummy.load_state_dict(torch.load('./torch_weights.pth', map_location = 'cpu')) 

but obviously, I'm getting errors due to the different layer names and sizes etc.

I've also tried:

class TorchModel(torch.nn.Module):
    def __init__(self, model) -> None:
        super().__init__()
        self.model = model.model
    
    def forward(self, image):
        return self.model([{"image":image}])[0]['instances']

But it doesn't work with .backbone, .layers etc

Answer 1

It looks like the notebook here focuses exactly on using GradCam for Faster-RCNN models, so the summary here should basically work for you. (Note that I have made some minor changes to this example to fit your variable names and add labels and bounding boxes.)

import cv2
from pytorch_grad_cam import AblationCAM, EigenCAM
from pytorch_grad_cam.ablation_layer import AblationLayerFasterRCNN
from pytorch_grad_cam.utils.model_targets import FasterRCNNBoxScoreTarget
from pytorch_grad_cam.utils.reshape_transforms import fasterrcnn_reshape_transform
from pytorch_grad_cam.utils.image import show_cam_on_image, scale_accross_batch_and_channels, scale_cam_image

# I have modified the code there slightly
im = cv2.imread("input.jpg")
outputs = model(im)
labels, boxes = outputs["instances"].pred_classes, outputs["instances"].pred_boxes

targets = [FasterRCNNBoxScoreTarget(labels=labels, bounding_boxes=boxes)]
target_layers = [model.model.backbone]
cam = AblationCAM(model.model,
                  target_layers, 
                  use_cuda=torch.nn.cuda.is_available(), 
                  reshape_transform=fasterrcnn_reshape_transform,
                  ablation_layer=AblationLayerFasterRCNN(),
                  ratio_channels_to_ablate=1.0)

# or a very fast alternative

cam = EigenCAM(model.model,
              target_layers, 
              use_cuda=torch.nn.cuda.is_available(), 
              reshape_transform=fasterrcnn_reshape_transform)