def create_fastercnn_backbone(backbone: str, fpn: bool = True, pretrained: str = None, trainable_backbone_layers: int = 3, **kwargs) -> nn.Module: """ Args: backbone (str): Supported backones are: "resnet18", "resnet34","resnet50", "resnet101", "resnet152", "resnext50_32x4d", "resnext101_32x8d", "wide_resnet50_2", "wide_resnet101_2", as resnets with fpn backbones. Without fpn backbones supported are: "resnet18", "resnet34", "resnet50","resnet101", "resnet152", "resnext101_32x8d", "mobilenet_v2", "vgg11", "vgg13", "vgg16", "vgg19", fpn (bool): If True then constructs fpn as well. pretrained (str): If None creates imagenet weights backbone. """ if fpn: # Creates a torchvision resnet model with fpn added. print("Resnet FPN Backbones works only for imagenet weights") backbone = resnet_fpn_backbone(backbone, pretrained=True, trainable_layers=trainable_backbone_layers, **kwargs) else: # This does not create fpn backbone, it is supported for all models print("FPN is not supported for Non Resnet Backbones") backbone, _ = create_torchvision_backbone(backbone, pretrained) return backbone
def __init__(self, pretrained_backbone: str, num_classes: int, learning_rate: float = 0.0001, pretrained: str = None, **kwargs,): super().__init__() self.num_classes = num_classes self.bottom, self.out_channels = create_torchvision_backbone(pretrained_backbone, pretrained) self.top = nn.Linear(self.out_channels, self.num_classes) self.learning_rate = learning_rate
def __init__(self, model_name: str, num_classes: int, pretrained: str = None): super().__init__() self.num_classes = num_classes self.bottom, self.out_channels = components.create_torchvision_backbone( model_name, pretrained) self.top = nn.Linear(self.out_channels, self.num_classes)