Faster R-CNN Add new Classifier (how to?)

Skylar_White: <p><div class="lightbox-wrapper"><a class="lightbox" href="https://discuss.pytorch.org/uploads/default/original/3X/8/5/85c5dcc22275d7592c308ae7d5bc12d5d48f08ed.png" data-download-href="https://discuss.pytorch.org/uploads/default/85c5dcc22275d7592c308ae7d5bc12d5d48f08ed" title="Planning for New Classifier"><img src="https://discuss.pytorch.org/uploads/default/optimized/3X/8/5/85c5dcc22275d7592c308ae7d5bc12d5d48f08ed_2_690x390.png" alt="Planning for New Classifier" data-base62-sha1="j5pq8Uii19jay78vnZrVfnnug3j" width="690" height="390" srcset="https://discuss.pytorch.org/uploads/default/optimized/3X/8/5/85c5dcc22275d7592c308ae7d5bc12d5d48f08ed_2_690x390.png, https://discuss.pytorch.org/uploads/default/optimized/3X/8/5/85c5dcc22275d7592c308ae7d5bc12d5d48f08ed_2_1035x585.png 1.5x, https://discuss.pytorch.org/uploads/default/original/3X/8/5/85c5dcc22275d7592c308ae7d5bc12d5d48f08ed.png 2x" data-dominant-color="EDEEEF"><div class="meta"><svg class="fa d-icon d-icon-far-image svg-icon" aria-hidden="true"><use href="#far-image"></use></svg><span class="filename">Planning for New Classifier</span><span class="informations">1200×679 143 KB</span><svg class="fa d-icon d-icon-discourse-expand svg-icon" aria-hidden="true"><use href="#discourse-expand"></use></svg></div></a></div></p>
<p>Hello guys, currently i have task to make object detection. Here I am trying to use Faster R-CNN, which is already pretrained, and later I will do fine tuning.</p>
<p>I also have another task, a (special) classification based on the input entered into the Faster R-CNN model, where my plan is to take the results of feature extraction from the “backbone”.</p>
<br>
sorry in advance if it's wrong, but i tried it like this.
<p>First, code to create a Faster R-CNN with a custom class.</p>
<pre><code class="lang-auto">def get_fasterrcnn(num_classes, pretrained=True):
    # load a model pre-trained on COCO
    if pretrained:
        w = torchvision.models.detection.FasterRCNN_ResNet50_FPN_V2_Weights.DEFAULT
        model = torchvision.models.detection.fasterrcnn_resnet50_fpn_v2(weights=w)
    else:
        model = torchvision.models.detection.fasterrcnn_resnet50_fpn_v2(weights=None)

    # get number of input features for the classifier
    in_features = model.roi_heads.box_predictor.cls_score.in_features

    # replace the pre-trained head with a new one
    model.roi_heads.box_predictor = FastRCNNPredictor(in_features, num_classes)

    return model
</code></pre>
<br>
<p>then next, I tried to take a feature map from the “backbone” layer (256? don’t know for sure) and I used it for my task custom classification for check fuel system.</p>
<pre><code class="lang-auto"># Create Custom Model
class CustomModelFuelSystem(nn.Module):
    def __init__(self, in_channels, out_channel=1):
        super().__init__()

        self.fuel_system_layers = nn.Sequential(
            nn.Linear(256, 128),
            nn.ReLU(),
            nn.Linear(128, out_channel)
        )

    def forward(self, x):
        return self.fuel_system_layers(x)
</code></pre>
<br>
after that I combined the Faster R-CNN model and the Custom model, like this
<pre><code class="lang-auto">class CustomFRCNN(nn.Module):
    def __init__(self, model, custom_model):
        super().__init__()

        self.faster_rcnn = model
        self.fuel_system = custom_model


    def forward(self, x):

        # do forward pass for model Faster R-CNN
        f_rcnn = self.faster_rcnn(x)

        # do forward pass for custom model layers
        ##  Get the output of the backbone and detach it
        backbone_output = f_rcnn[0]['backbone']

        fuel_sys = self.fuel_system(backbone_output)

        return f_rcnn, fuel_sys
</code></pre>
<br>
Next, I created these models for initialization
<pre><code class="lang-auto"># Create Model Faster R-CNN
model_rcnn = get_fasterrcnn(num_classes, pretrained=True)

# Create Model Custom Task
in_features_fuel = model_rcnn.backbone.out_channels
model_fuel = CustomModelFuelSystem(in_channels=in_features_fuel, out_channel=1)

# Create Model Custom FRCNN Task
model = CustomFRCNN(model=model_rcnn, custom_model=model_fuel)
</code></pre>
<br>
but when we try to classify with dummy data, there is an error like this:
<pre><code class="lang-auto"># Trying forward pass for dummies

# check sanity model
img_size = 800
x = torch.randn((1, 3, img_size, img_size))

model.eval()

test = model(x)

print(f"Input shape:\n {x.shape} \n")
print(test)
</code></pre>
<p><div class="lightbox-wrapper"><a class="lightbox" href="https://discuss.pytorch.org/uploads/default/original/3X/d/5/d512243c6d67c69face96f36bcddd7f51aea31d5.png" data-download-href="https://discuss.pytorch.org/uploads/default/d512243c6d67c69face96f36bcddd7f51aea31d5" title="error"><img src="https://discuss.pytorch.org/uploads/default/optimized/3X/d/5/d512243c6d67c69face96f36bcddd7f51aea31d5_2_690x325.png" alt="error" data-base62-sha1="uoUzilYeJTt7ntbWrNSLkgcng4l" width="690" height="325" srcset="https://discuss.pytorch.org/uploads/default/optimized/3X/d/5/d512243c6d67c69face96f36bcddd7f51aea31d5_2_690x325.png, https://discuss.pytorch.org/uploads/default/original/3X/d/5/d512243c6d67c69face96f36bcddd7f51aea31d5.png 1.5x, https://discuss.pytorch.org/uploads/default/original/3X/d/5/d512243c6d67c69face96f36bcddd7f51aea31d5.png 2x" data-dominant-color="3D3D3E"><div class="meta"><svg class="fa d-icon d-icon-far-image svg-icon" aria-hidden="true"><use href="#far-image"></use></svg><span class="filename">error</span><span class="informations">829×391 36 KB</span><svg class="fa d-icon d-icon-discourse-expand svg-icon" aria-hidden="true"><use href="#discourse-expand"></use></svg></div></a></div></p>
<p>I’m also confused, maybe someone can help?</p>
<p>thanks</p>