Save only the Backbone FPN from MaskRCNN for conversion to onnx format

Question

I have trained a MaskRCNN model successfully and now I am trying to save each of its stages (backbone with fpn, ROI classifier, Mask Head and RPN Head) in order to convert them to ONNX format. But it is not quite clear to me how I can decompose the model, currently, I have tried the following:

import torch.nn as nn
from torchvision.ops.feature_pyramid_network import FeaturePyramidNetwork, LastLevelMaxPool
from torchvision.models._utils import IntermediateLayerGetter
class BackboneWithFPN(nn.Module):
  def __init__(self):
    super(BackboneWithFPN, self).__init__()
    self.body = IntermediateLayerGetter()
    self.fpn = FeaturePyramidNetwork()
  def forward(self, x):
    x = self.body(x)
    return self.fpn(x)

Having already trained my model, my idea would be to call it through the methods to load it with those weights and start with the separation, something like this to later make the conversion to ONNX:

backbone_fpn = BackboneWithFPN()
backbone_fpn.body = model.backbone.body
backbone_fpn.fpn = model.backbone.fpn
backbone_fpn.eval()

However, I get several errors when I try to create this first stage of the model, any idea how I can solve this problem?

Save only the Backbone FPN from MaskRCNN for conversion to onnx format

0 Answers0