LastLevelP6P7 should allow custom in_feature for custom backbones
Summary:
LastLevelP6P7 had a hardcoded "res5" in feature key, PR adds parameter to overload this so RetinaNet FPN can be used with other backbones.
Pull Request resolved:
@@ -183,10 +183,10 @@ class LastLevelP6P7(nn.Module):
C5 feature.
"""
- def __init__(self, in_channels, out_channels):
+ def __init__(self, in_channels, out_channels, in_feature="res5"):
super().__init__()
self.num_levels = 2
- self.in_feature = "res5"
+ self.in_feature = in_feature
self.p6 = nn.Conv2d(in_channels, out_channels, 3, 2, 1)
self.p7 = nn.Conv2d(out_channels, out_channels, 3, 2, 1)
for module in [self.p6, self.p7]:
