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.
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”.
sorry in advance if it's wrong, but i tried it like this.
First, code to create a Faster R-CNN with a custom class.
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
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.
# 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)
after that I combined the Faster R-CNN model and the Custom model, like this
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
Next, I created these models for initialization
# 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)
but when we try to classify with dummy data, there is an error like this:
# 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)
I’m also confused, maybe someone can help?
thanks