Hello everyone,
I am currently trying to boost the classification accuracy of Faster-RCNN by using and auxiliary classification network for feature enrichment. To utilize the new inputs I need to add an additional argument in my box_predictor where x_cls is the last flatten layer of my auxiliary classification network.
class myRichDetector(nn.Module):
def __init__(self, in_channels, cls_flatten, num_classes):
super(myRichDetector, self).__init__()
# size of cls_flatten is equal to the size of the last flatten layer of the aux network
enriched_layer = in_channels + cls_flatten
self.pred_tor = FastRCNNPredictor(in_channels, num_classes)
self.pred_tor.cls_score = nn.Linear(enriched_layer, num_classes)
self.pred_tor.bbox_pred = nn.Linear(in_channels, num_classes * 4)
def forward(self, x, x_cls):
# enrich the output of the F-RCNN with the aux output
x_enr = torch.torch.cat((x, x_cls), 1)
if x.dim() == 4:
assert list(x.shape[2:]) == [1, 1]
x = x.flatten(start_dim=1)
scores = self.pred_tor.cls_score(x_enr)
bbox_deltas = self.pred_tor.bbox_pred(x)
return scores, bbox_deltas
class Model_detection(nn.Module):
def __init__(self, num_classes= None, in_features_cls=None, feature_extract=True, use_pretrained=True):
super(Model_detection, self).__init__()
self.model = models.detection.fasterrcnn_resnet50_fpn(pretrained=use_pretrained)
in_features = self.model.roi_heads.box_predictor.cls_score.in_features
# slice the model to extract the roi_heads and remove the box_predictor
self.model_my_roi_heads = nn.Sequential(*list(self.model.roi_heads.children())[:-1])
# remove the roi_head from the model
self.model = nn.Sequential(*list(self.model.children())[:-1])
# define the custom box_predictor
self.model_rich_head = myRichDetector(in_features, in_features_cls, num_classes)
self.set_parameter_requires_grad(self.model, feature_extract)
self.set_parameter_requires_grad(self.model_my_roi_heads, feature_extract)
self.set_parameter_requires_grad(self.model_rich_head, feature_extract)
def set_parameter_requires_grad(self, model, feature_extracting):
if feature_extracting:
for param in model.parameters():
param.requires_grad = False
def forward(self, x, target, cls_latent):
# normal F-RCNN forward pass
x = self.model(x, target)
x = self.model_my_roi_heads(x)
# introduce the new features for the prediction
x = self.model_rich_head(x, cls_latent)
return x
Now x = self.model(x, target)
gives me following error:
TypeError: forward() takes 2 positional arguments but 3 were given
Can anybody help me with this issue?
Thank you very much in advance