Wednesday, 5 October 2022

How to save Detectron2 model as a vanilla pytorch model?

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

I have my 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



from How to save Detectron2 model as a vanilla pytorch model?

No comments:

Post a Comment