-

features/convnext_features.py

@@ -40,76 +40,3 @@ def convnext_tiny_7_features(pretrained=False, **kwargs):
         model.classifier = nn.Identity()    
 
     return model
-
-class BasicGaussianMultiplierConv2D(nn.Conv2d):
-    def __init__(self, conv_layer, sigma=1.0, factor=50, device='cuda'):
-        super(BasicGaussianMultiplierConv2D, self).__init__(in_channels=conv_layer.in_channels, \
-                                         out_channels=conv_layer.out_channels, \
-                                         kernel_size=conv_layer.kernel_size, \
-                                         stride=conv_layer.stride, \
-                                         padding=conv_layer.padding, \
-                                         dilation=conv_layer.dilation, \
-                                         groups=conv_layer.groups, \
-                                         bias=conv_layer.bias is not None, \
-                                         padding_mode=conv_layer.padding_mode)
-
-        self.weight.data = conv_layer.weight.data.clone()
-        if conv_layer.bias is not None:
-            self.bias.data = conv_layer.bias.data.clone()
-        # self.gaussian_multiplier = self.generate_gaussian_kernel(size=self.weight.shape[-1], sigma=sigma).to(device)
-        # self.gaussian_multiplier.requires_grad = False
-        self.register_buffer('gaussian_multiplier', \
-                            self.generate_gaussian_kernel(size=self.weight.shape[-1], sigma=sigma).to(device))
-        self.factor = factor
-
-    def generate_gaussian_kernel(self, size, sigma=1.0):
-        kernel = np.fromfunction(
-            lambda x, y: (1 / (2 * np.pi * sigma**2)) * np.exp(-((x - (size-1)/2)**2 + (y - (size-1)/2)**2) / (2 * sigma**2)),
-            (size, size)
-        )
-        kernel = torch.Tensor(kernel)
-        kernel /= kernel.sum()
-        kernel = kernel
-        kernel = kernel.expand(1, 1, size, size)
-        return kernel
-
-    def forward(self, input):
-        if self.bias is not None:
-            bias = self.bias.data
-        else:
-            bias = None
-        gaussian_multiplied_kernel = self.weight.data * self.gaussian_multiplier * self.factor
-        return self._conv_forward(input, gaussian_multiplied_kernel, bias)
-
-def replace_CNBlock_conv_layer_with_gaussian_multiplier(model: torchvision.models.convnext.CNBlock, sigma=1.0, factor=50, device='cuda'):
-    first_sequential_inside_cnblock = list(model.children())[0]
-    cnblock_7x7_layer = first_sequential_inside_cnblock[0]
-    first_sequential_inside_cnblock[0] = BasicGaussianMultiplierConv2D(cnblock_7x7_layer, sigma=sigma, factor=factor, device=device)
-
-def replace_convnext_convlayers_with_gaussian_multipliers(model, sigma=1.0, factor=50, device='cuda'):
-    for n, module in model.named_children():
-        if isinstance(module, torchvision.models.convnext.CNBlock):
-            replace_CNBlock_conv_layer_with_gaussian_multiplier(module, sigma=sigma, factor=factor, device=device)
-        elif len(list(module.children())) > 0:
-            replace_convlayers_with_gaussian_convnext(module, sigma=sigma, factor=factor, device=device)
-    return model
-
-def apply_gaussian_multiplier_to_convnext_stage(original_model, stage, sigma=1.0, factor=50, device='cuda'):
-    if stage == 1:
-        print('stage', stage)
-        for child in list(original_model.features.children())[:2]:
-            replace_convnext_convlayers_with_gaussian_multipliers(child, sigma=sigma, factor=factor, device=device)
-    elif stage == 2:
-        print('stage', stage)
-        for child in list(original_model.features.children())[2:4]:
-            replace_convnext_convlayers_with_gaussian_multipliers(child, sigma=sigma, factor=factor, device=device)
-    elif stage == 3:
-        print('stage', stage)
-        for child in list(original_model.features.children())[4:6]:
-            replace_convnext_convlayers_with_gaussian_multipliers(child, sigma=sigma, factor=factor, device=device)
-    elif stage == 4:
-        print('stage', stage)
-        for child in list(original_model.features.children())[6:8]:
-            replace_convnext_convlayers_with_gaussian_multipliers(child, sigma=sigma, factor=factor, device=device)
-
-    return original_model

hcompnet/model.py

@@ -2,7 +2,7 @@
 import torch
 import torch.nn as nn
 import torch.nn.functional as F
-from features.convnext_features import convnext_tiny_26_features, convnext_tiny_13_features, convnext_tiny_7_features, apply_gaussian_multiplier_to_convnext_stage
+from features.convnext_features import convnext_tiny_26_features, convnext_tiny_13_features, convnext_tiny_7_features
 import torch
 from torch import Tensor
 from util.node import Node
@@ -87,13 +87,7 @@ def get_classification_layers(self):
         return [getattr(self, attr) for attr in dir(self) if attr.endswith('_classification')] 
 
 
-base_architecture_to_features = {'resnet18': resnet18_features,
-                                 'resnet34': resnet34_features,
-                                 'resnet50': resnet50_features,
-                                 'resnet50_inat': resnet50_features_inat,
-                                 'resnet101': resnet101_features,
-                                 'resnet152': resnet152_features,
-                                 'convnext_tiny_26': convnext_tiny_26_features,
+base_architecture_to_features = {'convnext_tiny_26': convnext_tiny_26_features,
                                  'convnext_tiny_13': convnext_tiny_13_features,
                                  'convnext_tiny_7': convnext_tiny_7_features}
 

util/args.py

@@ -25,7 +25,7 @@ def get_args() -> argparse.Namespace:
     parser.add_argument('--net',
                         type=str,
                         default='convnext_tiny_26',
-                        help='Base network used as backbone of HComP-Net. Default is convnext_tiny_26 with adapted strides to output 26x26 latent representations. Other option is convnext_tiny_13 that outputs 13x13 (smaller and faster to train, less fine-grained). Pretrained network on iNaturalist is only available for resnet50_inat. Options are: resnet18, resnet34, resnet50, resnet50_inat, resnet101, resnet152, convnext_tiny_26 and convnext_tiny_13.')
+                        help='Base network used as backbone of HComP-Net. Default is convnext_tiny_26 with adapted strides to output 26x26 latent representations. Other option is convnext_tiny_13 that outputs 13x13 (smaller and faster to train, less fine-grained)')
     parser.add_argument('--batch_size',
                         type=int,
                         default=64,
@@ -237,43 +237,7 @@ def get_optimizer_nn(net, args: argparse.Namespace) -> torch.optim.Optimizer:
     params_to_train = []
     params_backbone = []
     # set up optimizer
-    if 'resnet50' in args.net: 
-        # freeze resnet50 except last convolutional layer
-        for name,param in net.module._net.named_parameters():
-            if 'layer4.2' in name:
-                params_to_train.append(param)
-            elif 'layer4' in name or 'layer3' in name:
-                params_to_freeze.append(param)
-            elif 'layer2' in name:
-                params_backbone.append(param)
-            else: #such that model training fits on one gpu. 
-                param.requires_grad = False
-                # params_backbone.append(param)
-    elif 'resnet18' in args.net: 
-        # freeze resnet50 except last convolutional layer
-        for name,param in net.module._net.named_parameters():
-            if 'layer4.1' in name:
-                params_to_train.append(param)
-            elif 'layer4' in name or 'layer3' in name:
-                params_to_freeze.append(param)
-            # elif 'layer2' in name:
-            #     params_backbone.append(param)
-            else: #such that model training fits on one gpu. 
-                # param.requires_grad = False
-                params_backbone.append(param)
-    elif 'resnet34' in args.net: 
-        # freeze resnet50 except last convolutional layer
-        for name,param in net.module._net.named_parameters():
-            if 'layer4.2' in name:
-                params_to_train.append(param)
-            elif 'layer4' in name or 'layer3' in name:
-                params_to_freeze.append(param)
-            # elif 'layer2' in name:
-            #     params_backbone.append(param)
-            else: #such that model training fits on one gpu. 
-                # param.requires_grad = False
-                params_backbone.append(param)
-    elif 'convnext' in args.net:
+    if 'convnext' in args.net:
         print("chosen network is convnext", flush=True)
         for name,param in net.module._net.named_parameters():
             if 'features.7.2' in name: