Hi im trying to train a SSD With mobilenet backbone
I want to delete the output layer of backbone and put 2 convolutional output layers one for classification and another for location
class ssdLightModel(torch.nn.Module):
def __init__(self):
super(ssdLightModel, self).__init__()
self.backbone = models.mobilenet_v2(weights='MobileNet_V2_Weights.IMAGENET1K_V1')
infeatures = self.backbone._modules['classifier'][-1].in_features
labelOutfeatures = len(SCALES) * len(ASPECT_RATIOS) * 2
boxOutfeatures = len(SCALES) * len(ASPECT_RATIOS) * 4
self.backbone = self.backbone.features
self.fc1 = nn.Conv2d(in_channels=infeatures, out_channels=labelOutfeatures, kernel_size=1, bias=True)
self.fc2 = nn.Conv2d(in_channels=infeatures, out_channels=boxOutfeatures, kernel_size=1, bias=True)
# kernel_size=(FMAP_ROW, FMAP_COL)
def forward(self, x):
x = self.backbone(x)
x1 = self.fc1(x)
x2 = self.fc2(x)
return x1, x2
This is my model definition, Did i do something wrong? Because my model learned nothing….