def build_model()

in src/similarity/cnn.py [0:0]

• 20 lines of code
• 8 McCabe index (conditional complexity)

    def build_model(self, dims, v):
        
        if v == 18:
            cnn = models.resnet18(pretrained=True)
        elif v == 34:
            cnn = models.resnet34(pretrained=True)
        elif v == 50:
            cnn = models.resnet50(pretrained=True)
        elif v == 101:
            cnn = models.resnet101(pretrained=True)
        elif v == 152:
            cnn = models.resnet152(pretrained=True)
        else:
            cnn = models.resnet152(pretrained=True)
            
        lastlayer_in = cnn.fc.in_features
        cnn.fc = nn.Linear(lastlayer_in, dims)
        
        active_layers = {"fc.weight":1,"fc.bias":1}
        for name, param in cnn.named_parameters():
            if name not in active_layers:
                param.requires_grad = False
         
        self.sim_model = cnn