Hellow there.
I’m currently trying to build a custom object detection model based on Faster R-CNN.
The structure of the custom CNN model is such that the final output is the class of the object, the boundary box, and the attributes of the object .
So, I wrote the following code.
class CustomPredictor(nn.Module):
def __init__(self,in_channels,num_classes):
super(CustomPredictor,self).__init__()
self.cls_score = nn.Linear(in_channels, num_classes)
self.bbox_pred = nn.Linear(in_channels, num_classes * 4)
self.additional_layer = nn.Linear(in_channels,1) #this is the additional layer
def forward(self,x):
if x.dim() == 4:
assert list(x.shape[2:]) == [1, 1]
x = x.flatten(start_dim=1)
pred_leaf_age = self.additional_layer(x)
scores = self.cls_score(x)
bbox_deltas = self.bbox_pred(x)
return scores, bbox_deltas, pred_leaf_age
class CustomModel(nn.Module):
def __init__(self):
super(CustomModel, self).__init__()
self.model = torchvision.models.detection.fasterrcnn_resnet50_fpn(pretrained=True)
self.in_features = self.model.roi_heads.box_head.fc7.out_features
self.model.roi_heads.box_predictor = CustomPredictor(self.in_features,6)
def forward(self, x):
The problem here is how to define the foward of the CustomModel class.
Could you give me any advice?
Best regards.