Upload all files

DIC.py

@@ -0,0 +1,17 @@
+import torch
+from pathlib import Path
+
+
+dir=Path.home() / f"tmp/resnet50/CUB2011/123456/"
+dic=torch.load(dir/ f"SlDD_Selection_50.pt")
+
+print (dic)
+
+#if 'linear.selection' in dic.keys():
+    #print("key 'linear.selection' exist")
+#else:
+    #print("no such key")
+
+
+
+

FeatureDiversityLoss.py

@@ -0,0 +1,59 @@
+import torch
+from torch import nn
+
+"""
+Feature Diversity Loss:
+Usage to replicate paper:
+Call 
+loss_function = FeatureDiversityLoss(0.196, linear) 
+to inititalize loss with linear layer of model.
+At each mini batch get feature maps (Output of final convolutional layer) and add to Loss:
+loss += loss_function(feature_maps, outputs)
+"""
+
+
+class FeatureDiversityLoss(nn.Module):
+    def __init__(self, scaling_factor, linear):
+        super().__init__()
+        self.scaling_factor = scaling_factor #* 0
+        print("Scaling Factor: ", self.scaling_factor)
+        self.linearLayer = linear
+
+    def initialize(self, linearLayer):
+        self.linearLayer = linearLayer
+
+    def get_weights(self, outputs):
+        weight_matrix = self.linearLayer.weight
+        weight_matrix = torch.abs(weight_matrix)
+        top_classes = torch.argmax(outputs, dim=1)
+        relevant_weights = weight_matrix[top_classes]
+        return relevant_weights
+
+    def forward(self, feature_maps, outputs):
+        relevant_weights = self.get_weights(outputs)
+        relevant_weights = norm_vector(relevant_weights)
+        feature_maps = preserve_avg_func(feature_maps)
+        flattened_feature_maps = feature_maps.flatten(2)
+        batch, features, map_size = flattened_feature_maps.size()
+        relevant_feature_maps = flattened_feature_maps * relevant_weights[..., None]
+        diversity_loss = torch.sum(
+            torch.amax(relevant_feature_maps, dim=1))
+        return -diversity_loss / batch * self.scaling_factor
+
+
+def norm_vector(x):
+    return x / (torch.norm(x, dim=1) + 1e-5)[:, None]
+
+
+def preserve_avg_func(x):
+    avgs = torch.mean(x, dim=[2, 3])
+    max_avgs = torch.max(avgs, dim=1)[0]
+    scaling_factor = avgs / torch.clamp(max_avgs[..., None], min=1e-6)
+    softmaxed_maps = softmax_feature_maps(x)
+    scaled_maps = softmaxed_maps * scaling_factor[..., None, None]
+    return scaled_maps
+
+
+def softmax_feature_maps(x):
+    return torch.softmax(x.reshape(x.size(0), x.size(1), -1), 2).view_as(x)
+

app.py

@@ -0,0 +1,143 @@
+import gradio as gr
+from load_model import extract_sel_mean_std_bias_assignemnt
+from pathlib import Path
+from architectures.model_mapping import get_model
+from configs.dataset_params import dataset_constants
+import torch
+import torchvision.transforms as transforms
+import pandas as pd
+import cv2
+import numpy as np
+
+def overlapping_features_on_input(model,output, feature_maps, input, target):
+    W=model.linear.layer.weight
+    output=output.detach().cpu().numpy()
+    feature_maps=feature_maps.detach().cpu().numpy().squeeze()
+
+    if target !=None:
+     label=target
+    else:
+     label=np.argmax(output)+1
+
+    Interpretable_Selection= W[label,:]
+    print("W",Interpretable_Selection)
+    input_np=np.array(input)
+    h,w= input.shape[:2]
+    print("h,w:",h,w)
+    Interpretable_Features=[]
+    Feature_image_list=[]
+    for S in range(len(Interpretable_Selection)):
+        if Interpretable_Selection[S] > 0:
+               Interpretable_Features.append(feature_maps[S])
+               Feature_image=cv2.resize(feature_maps[S],(w,h))
+               Feature_image=((Feature_image-np.min(Feature_image))/(np.max(Feature_image)-np.min(Feature_image)))*255
+               Feature_image=Feature_image.astype(np.uint8)
+               Feature_image=cv2.applyColorMap(Feature_image,cv2.COLORMAP_JET)
+               Feature_image=0.3*Feature_image+0.7*input_np
+               Feature_image=np.clip(Feature_image, 0, 255).astype(np.uint8)
+               Feature_image_list.append(Feature_image)
+               #path_to_featureimage=f"/home/qixuan/tmp/FeatureImage/FI{S}.jpg"
+               #cv2.imwrite(path_to_featureimage,Feature_image)
+    print("len of Features:",len(Interpretable_Features))
+
+    return Feature_image_list
+
+
+def genreate_intepriable_output(input,dataset="CUB2011", arch="resnet50",seed=123456, model_type="qsenn", n_features = 50, n_per_class=5, img_size=448, reduced_strides=False, folder = None):
+    n_classes = dataset_constants[dataset]["num_classes"]
+  
+    model = get_model(arch, n_classes, reduced_strides)
+    tr=transforms.ToTensor()
+    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+    if folder is None:
+        folder = Path(f"tmp/{arch}/{dataset}/{seed}/")
+    
+    state_dict = torch.load(folder / f"{model_type}_{n_features}_{n_per_class}_FinetunedModel.pth")
+    selection= torch.load(folder / f"SlDD_Selection_50.pt")
+    state_dict['linear.selection']=selection
+    
+    feature_sel, sparse_layer, current_mean, current_std, bias_sparse = extract_sel_mean_std_bias_assignemnt(state_dict)
+    model.set_model_sldd(feature_sel, sparse_layer, current_mean, current_std, bias_sparse)
+    model.load_state_dict(state_dict)
+
+    input = tr(input)
+    input= input.unsqueeze(0)
+    input= input.to(device)
+    model = model.to(device)
+    output, feature_maps, final_features = model(input, with_feature_maps=True, with_final_features=True)
+    print("final features:",final_features)
+    output=output.detach().cpu().numpy()
+    output= np.argmax(output)+1
+    
+
+    print("outputclass:",output)
+    data_dir=Path("tmp/Datasets/CUB200/CUB_200_2011/")
+    labels = pd.read_csv(data_dir/"image_class_labels.txt", sep=' ', names=['img_id', 'target'])
+    namelist=pd.read_csv(data_dir/"images.txt",sep=' ',names=['img_id','file_name'])
+    classlist=pd.read_csv(data_dir/"classes.txt",sep=' ',names=['cl_id','class_name'])
+    options_output=labels[labels['target']==output]
+    options_output=options_output.sample(1)
+    others=labels[labels['target']!=output]
+    options_others=others.sample(3)
+    options = pd.concat([options_others, options_output], ignore_index=True)
+    shuffled_options = options.sample(frac=1).reset_index(drop=True)
+    print("shuffled:",shuffled_options)
+    op=[]
+
+    for i in shuffled_options['img_id']:
+        print(i)
+        filenames=namelist.loc[namelist['img_id']==i,'file_name'].values[0]
+        targets=shuffled_options.loc[shuffled_options['img_id']==i,'target'].values[0]
+        print("targets",targets)
+        print("name",filenames)
+
+        classes=classlist.loc[classlist['cl_id']==targets, 'class_name'].values[0]
+        print(data_dir/f"images/{filenames}")
+
+        op_img=cv2.imread(data_dir/f"images/{filenames}")
+
+        op_images=tr(op_img)
+        op_images=op_images.unsqueeze(0)
+        op_images=op_images.to(device)
+        OP, feature_maps_op =model(op_images,with_feature_maps=True,with_final_features=False)
+        print("OP:",OP,
+              "feature_maps_op:",feature_maps_op.shape)
+        opt= overlapping_features_on_input(model,OP, feature_maps_op,op_img,targets)
+        op+=opt
+    
+    return op
+
+def post_next_image(op):
+    if len(op)<=1:
+        return [],None, "all done, thank you!"
+    else:
+        op=op[1:len(op)]
+        return op,op[0], "Is this feature also in your input?"
+
+def get_features_on_interface(input):
+    op=genreate_intepriable_output(input,dataset="CUB2011", 
+                                arch="resnet50",seed=123456, 
+                                model_type="qsenn", n_features = 50,n_per_class=5,
+                                img_size=448, reduced_strides=False, folder = None)
+    return op, op[0],"Is this feature also in your input?",gr.update(interactive=False)
+
+
+with gr.Blocks() as demo:
+
+    gr.Markdown("<h1 style='text-align: center;'>Interiable Bird Classification</h1>")
+    image_input=gr.Image()
+    image_output=gr.Image()
+    text_output=gr.Markdown()
+    but_generate=gr.Button("Get some interpriable Features")
+    but_feedback_y=gr.Button("Yes")
+    but_feedback_n=gr.Button("No")
+    image_list = gr.State([])
+    but_generate.click(fn=get_features_on_interface, inputs=image_input, outputs=[image_list,image_output,text_output,but_generate])
+    but_feedback_y.click(fn=post_next_image, inputs=image_list, outputs=[image_list,image_output,text_output])
+    but_feedback_n.click(fn=post_next_image, inputs=image_list, outputs=[image_list,image_output,text_output])
+
+demo.launch()  
+  
+    
+
+    

architectures/FinalLayer.py

@@ -0,0 +1,36 @@
+import torch
+from torch import nn
+
+from architectures.SLDDLevel import SLDDLevel
+
+
+class FinalLayer():
+    def __init__(self, num_classes,  n_features):
+        super().__init__()
+        self.avgpool = torch.nn.AdaptiveAvgPool2d((1, 1))
+        self.linear = nn.Linear(n_features, num_classes)
+        self.featureDropout = torch.nn.Dropout(0.2)
+        self.selection = None
+
+    def transform_output(self,  feature_maps, with_feature_maps,
+                         with_final_features):
+        if self.selection is not None:
+            feature_maps = feature_maps[:, self.selection]
+        x = self.avgpool(feature_maps)
+        pre_out = torch.flatten(x, 1)
+        final_features = self.featureDropout(pre_out)
+        final = self.linear(final_features)
+        final = [final]
+        if with_feature_maps:
+            final.append(feature_maps)
+        if with_final_features:
+            final.append(final_features)
+        if len(final) == 1:
+            final = final[0]
+        return final
+
+
+    def set_model_sldd(self, selection, weight, mean, std, bias = None):
+        self.selection = selection
+        self.linear = SLDDLevel(selection, weight, mean, std, bias)
+        self.featureDropout = torch.nn.Dropout(0.1)

architectures/SLDDLevel.py

@@ -0,0 +1,37 @@
+import torch.nn
+
+
+class SLDDLevel(torch.nn.Module):
+    def __init__(self, selection, weight_at_selection,mean, std, bias=None):
+        super().__init__()
+        self.register_buffer('selection', torch.tensor(selection, dtype=torch.long))
+        num_classes,        n_features = weight_at_selection.shape
+        selected_mean = mean
+        selected_std = std
+        if len(selected_mean) != len(selection):
+            selected_mean = selected_mean[selection]
+            selected_std = selected_std[selection]
+        self.mean = torch.nn.Parameter(selected_mean)
+        self.std = torch.nn.Parameter(selected_std)
+        if bias is not None:
+            self.layer = torch.nn.Linear(n_features, num_classes)
+            self.layer.bias = torch.nn.Parameter(bias, requires_grad=False)
+        else:
+            self.layer = torch.nn.Linear(n_features, num_classes, bias=False)
+        self.layer.weight = torch.nn.Parameter(weight_at_selection, requires_grad=False)
+
+    @property
+    def weight(self):
+        return self.layer.weight
+
+    @property
+    def bias(self):
+        if self.layer.bias is None:
+            return torch.zeros(self.layer.out_features)
+        else:
+            return self.layer.bias
+
+
+    def forward(self, input):
+        input = (input - self.mean) / torch.clamp(self.std, min=1e-6)
+        return self.layer(input)

architectures/model_mapping.py

@@ -0,0 +1,7 @@
+from architectures.resnet import resnet50
+
+
+def get_model(arch, num_classes, changed_strides=True):
+    if arch == "resnet50":
+        model = resnet50(True, num_classes=num_classes, changed_strides=changed_strides)
+    return  model

architectures/resnet.py

@@ -0,0 +1,420 @@
+import copy
+import time
+
+import torch
+import torch.nn as nn
+from torch.hub import load_state_dict_from_url
+from torchvision.models import get_model
+
+# from scripts.modelExtensions.crossModelfunctions import init_experiment_stuff
+
+
+
+__all__ = ['ResNet', 'resnet18', 'resnet34', 'resnet50', 'resnet101',
+           'resnet152', 'resnext50_32x4d', 'resnext101_32x8d',
+           'wide_resnet50_2', 'wide_resnet101_2',
+           'wide_resnet50_3', 'wide_resnet50_4', 'wide_resnet50_5',
+           'wide_resnet50_6', ]
+
+from architectures.FinalLayer import FinalLayer
+from architectures.utils import SequentialWithArgs
+
+model_urls = {
+    'resnet18': 'https://download.pytorch.org/models/resnet18-5c106cde.pth',
+    'resnet34': 'https://download.pytorch.org/models/resnet34-333f7ec4.pth',
+    'resnet50': 'https://download.pytorch.org/models/resnet50-19c8e357.pth',
+    'resnet101': 'https://download.pytorch.org/models/resnet101-5d3b4d8f.pth',
+    'resnet152': 'https://download.pytorch.org/models/resnet152-b121ed2d.pth',
+    'resnext50_32x4d': 'https://download.pytorch.org/models/resnext50_32x4d-7cdf4587.pth',
+    'resnext101_32x8d': 'https://download.pytorch.org/models/resnext101_32x8d-8ba56ff5.pth',
+    'wide_resnet50_2': 'https://download.pytorch.org/models/wide_resnet50_2-95faca4d.pth',
+    'wide_resnet101_2': 'https://download.pytorch.org/models/wide_resnet101_2-32ee1156.pth',
+}
+
+
+def conv3x3(in_planes, out_planes, stride=1, groups=1, dilation=1):
+    """3x3 convolution with padding"""
+    return nn.Conv2d(in_planes, out_planes, kernel_size=3, stride=stride,
+                     padding=dilation, groups=groups, bias=False, dilation=dilation)
+
+
+def conv1x1(in_planes, out_planes, stride=1):
+    """1x1 convolution"""
+    return nn.Conv2d(in_planes, out_planes, kernel_size=1, stride=stride, bias=False)
+
+
+class BasicBlock(nn.Module):
+    expansion = 1
+    __constants__ = ['downsample']
+
+    def __init__(self, inplanes, planes, stride=1, downsample=None, groups=1,
+                 base_width=64, dilation=1, norm_layer=None, features=None):
+        super(BasicBlock, self).__init__()
+        if norm_layer is None:
+            norm_layer = nn.BatchNorm2d
+        if groups != 1 or base_width != 64:
+            raise ValueError('BasicBlock only supports groups=1 and base_width=64')
+        if dilation > 1:
+            raise NotImplementedError("Dilation > 1 not supported in BasicBlock")
+        # Both self.conv1 and self.downsample layers downsample the input when stride != 1
+        self.conv1 = conv3x3(inplanes, planes, stride)
+        self.bn1 = norm_layer(planes)
+        self.relu = nn.ReLU(inplace=True)
+        self.conv2 = conv3x3(planes, planes)
+        self.bn2 = norm_layer(planes)
+        self.downsample = downsample
+        self.stride = stride
+
+
+    def forward(self, x,  no_relu=False):
+        identity = x
+
+        out = self.conv1(x)
+        out = self.bn1(out)
+        out = self.relu(out)
+
+        out = self.conv2(out)
+        out = self.bn2(out)
+
+        if self.downsample is not None:
+            identity = self.downsample(x)
+
+
+
+        out += identity
+
+        if no_relu:
+            return out
+        return self.relu(out)
+
+
+class Bottleneck(nn.Module):
+    expansion = 4
+    __constants__ = ['downsample']
+
+    def __init__(self, inplanes, planes, stride=1, downsample=None, groups=1,
+                 base_width=64, dilation=1, norm_layer=None, features=None):
+        super(Bottleneck, self).__init__()
+        if norm_layer is None:
+            norm_layer = nn.BatchNorm2d
+        width = int(planes * (base_width / 64.)) * groups
+        # Both self.conv2 and self.downsample layers downsample the input when stride != 1
+        self.conv1 = conv1x1(inplanes, width)
+        self.bn1 = norm_layer(width)
+        self.conv2 = conv3x3(width, width, stride, groups, dilation)
+        self.bn2 = norm_layer(width)
+        if features is None:
+            self.conv3 = conv1x1(width, planes * self.expansion)
+            self.bn3 = norm_layer(planes * self.expansion)
+        else:
+            self.conv3 = conv1x1(width, features)
+            self.bn3 = norm_layer(features)
+
+        self.relu = nn.ReLU(inplace=True)
+        self.downsample = downsample
+        self.stride = stride
+
+    def forward(self, x,  no_relu=False, early_exit=False):
+        identity = x
+        out = self.conv1(x)
+        out = self.bn1(out)
+        out = self.relu(out)
+
+        out = self.conv2(out)
+        out = self.bn2(out)
+        out = self.relu(out)
+
+        out = self.conv3(out)
+        out = self.bn3(out)
+
+        if self.downsample is not None:
+            identity = self.downsample(x)
+
+        out += identity
+
+        if no_relu:
+            return out
+        return self.relu(out)
+
+
+class ResNet(nn.Module, FinalLayer):
+
+    def __init__(self, block, layers, num_classes=1000, zero_init_residual=False,
+                 groups=1, width_per_group=64, replace_stride_with_dilation=None,
+                 norm_layer=None, changed_strides=False,):
+        super(ResNet, self).__init__()
+        if norm_layer is None:
+            norm_layer = nn.BatchNorm2d
+        self._norm_layer = norm_layer
+        widths = [64, 128, 256, 512]
+        self.inplanes = 64
+        self.dilation = 1
+        if replace_stride_with_dilation is None:
+            # each element in the tuple indicates if we should replace
+            # the 2x2 stride with a dilated convolution instead
+            replace_stride_with_dilation = [False, False, False]
+        if len(replace_stride_with_dilation) != 3:
+            raise ValueError("replace_stride_with_dilation should be None "
+                             "or a 3-element tuple, got {}".format(replace_stride_with_dilation))
+        self.groups = groups
+        self.base_width = width_per_group
+        self.conv1 = nn.Conv2d(3, self.inplanes, kernel_size=7, stride=2, padding=3,
+                               bias=False)
+        self.bn1 = norm_layer(self.inplanes)
+        self.relu = nn.ReLU(inplace=True)
+        self.maxpool = nn.MaxPool2d(kernel_size=3, stride=2, padding=1)
+        self.layer1 = self._make_layer(block, 64, layers[0])
+        self.layer2 = self._make_layer(block, 128, layers[1], stride=2,
+                                       dilate=replace_stride_with_dilation[0])
+        self.sstride = 2
+        if changed_strides:
+            self.sstride = 1
+        self.layer3 = self._make_layer(block, 256, layers[2], stride=self.sstride,
+                                       dilate=replace_stride_with_dilation[1])
+        self.stride = 2
+
+        if changed_strides:
+            self.stride = 1
+        self.layer4 = self._make_layer(block, 512, layers[3], stride=self.stride,
+                                       dilate=replace_stride_with_dilation[2])
+        FinalLayer.__init__(self, num_classes, 512 * block.expansion)
+        for m in self.modules():
+            if isinstance(m, nn.Conv2d):
+                nn.init.kaiming_normal_(m.weight, mode='fan_out', nonlinearity='relu')
+            elif isinstance(m, (nn.BatchNorm2d, nn.GroupNorm)):
+                nn.init.constant_(m.weight, 1)
+                nn.init.constant_(m.bias, 0)
+
+        # Zero-initialize the last BN in each residual branch,
+        # so that the residual branch starts with zeros, and each residual block behaves like an identity.
+        # This improves the model by 0.2~0.3% according to https://arxiv.org/abs/1706.02677
+        if zero_init_residual:
+            for m in self.modules():
+                if isinstance(m, Bottleneck):
+                    nn.init.constant_(m.bn3.weight, 0)
+                elif isinstance(m, BasicBlock):
+                    nn.init.constant_(m.bn2.weight, 0)
+
+    def _make_layer(self, block, planes, blocks, stride=1, dilate=False, last_block_f=None):
+        norm_layer = self._norm_layer
+        downsample = None
+        previous_dilation = self.dilation
+        if dilate:
+            self.dilation *= stride
+            stride = 1
+        if stride != 1 or self.inplanes != planes * block.expansion:
+            downsample = nn.Sequential(
+                conv1x1(self.inplanes, planes * block.expansion, stride),
+                norm_layer(planes * block.expansion),
+            )
+
+        layers = []
+        layers.append(block(self.inplanes, planes, stride, downsample, self.groups,
+                            self.base_width, previous_dilation, norm_layer))
+        self.inplanes = planes * block.expansion
+        for _ in range(1, blocks):
+            krepeep = None
+            if last_block_f is not None and _ == blocks - 1:
+                krepeep = last_block_f
+            layers.append(block(self.inplanes, planes, groups=self.groups,
+                                base_width=self.base_width, dilation=self.dilation,
+                                norm_layer=norm_layer, features=krepeep))
+
+        return SequentialWithArgs(*layers)
+
+    def _forward(self, x, with_feature_maps=False, with_final_features=False):
+        x = self.conv1(x)
+        x = self.bn1(x)
+        x = self.relu(x)
+        x = self.maxpool(x)
+
+        x = self.layer1(x)
+        x = self.layer2(x)
+        x = self.layer3(x)
+        feature_maps = self.layer4(x,  no_relu=True)
+        feature_maps = torch.functional.F.relu(feature_maps)
+        return self.transform_output( feature_maps,  with_feature_maps,
+                                     with_final_features)
+
+    # Allow for accessing forward method in a inherited class
+    forward = _forward
+
+
+def _resnet(arch, block, layers, pretrained, progress, **kwargs):
+    model = ResNet(block, layers, **kwargs)
+    if pretrained:
+        state_dict = load_state_dict_from_url(model_urls[arch],
+                                                      progress=progress)
+        if kwargs["num_classes"] == 1000:
+            state_dict["linear.weight"] = state_dict["fc.weight"]
+            state_dict["linear.bias"] = state_dict["fc.bias"]
+        model.load_state_dict(state_dict, strict=False)
+    return model
+
+
+def resnet18(pretrained=False, progress=True, **kwargs):
+    r"""ResNet-18 model from
+    `"Deep Residual Learning for Image Recognition" <https://arxiv.org/pdf/1512.03385.pdf>`_
+
+    Args:
+        pretrained (bool): If True, returns a model pre-trained on ImageNet
+        progress (bool): If True, displays a progress bar of the download to stderr
+    """
+    return _resnet('resnet18', BasicBlock, [2, 2, 2, 2], pretrained, progress,
+                   **kwargs)
+
+
+def resnet34(pretrained=False, progress=True, **kwargs):
+    r"""ResNet-34 model from
+    `"Deep Residual Learning for Image Recognition" <https://arxiv.org/pdf/1512.03385.pdf>`_
+
+    Args:
+        pretrained (bool): If True, returns a model pre-trained on ImageNet
+        progress (bool): If True, displays a progress bar of the download to stderr
+    """
+    return _resnet('resnet34', BasicBlock, [3, 4, 6, 3], pretrained, progress,
+                   **kwargs)
+
+
+def resnet50(pretrained=False, progress=True, **kwargs):
+    r"""ResNet-50 model from
+    `"Deep Residual Learning for Image Recognition" <https://arxiv.org/pdf/1512.03385.pdf>`_
+
+    Args:
+        pretrained (bool): If True, returns a model pre-trained on ImageNet
+        progress (bool): If True, displays a progress bar of the download to stderr
+    """
+    return _resnet('resnet50', Bottleneck, [3, 4, 6, 3], pretrained, progress,
+                   **kwargs)
+
+
+def resnet101(pretrained=False, progress=True, **kwargs):
+    r"""ResNet-101 model from
+    `"Deep Residual Learning for Image Recognition" <https://arxiv.org/pdf/1512.03385.pdf>`_
+
+    Args:
+        pretrained (bool): If True, returns a model pre-trained on ImageNet
+        progress (bool): If True, displays a progress bar of the download to stderr
+    """
+    return _resnet('resnet101', Bottleneck, [3, 4, 23, 3], pretrained, progress,
+                   **kwargs)
+
+
+def resnet152(pretrained=False, progress=True, **kwargs):
+    r"""ResNet-152 model from
+    `"Deep Residual Learning for Image Recognition" <https://arxiv.org/pdf/1512.03385.pdf>`_
+
+    Args:
+        pretrained (bool): If True, returns a model pre-trained on ImageNet
+        progress (bool): If True, displays a progress bar of the download to stderr
+    """
+    return _resnet('resnet152', Bottleneck, [3, 8, 36, 3], pretrained, progress,
+                   **kwargs)
+
+
+def resnext50_32x4d(pretrained=False, progress=True, **kwargs):
+    r"""ResNeXt-50 32x4d model from
+    `"Aggregated Residual Transformation for Deep Neural Networks" <https://arxiv.org/pdf/1611.05431.pdf>`_
+
+    Args:
+        pretrained (bool): If True, returns a model pre-trained on ImageNet
+        progress (bool): If True, displays a progress bar of the download to stderr
+    """
+    kwargs['groups'] = 32
+    kwargs['width_per_group'] = 4
+    return _resnet('resnext50_32x4d', Bottleneck, [3, 4, 6, 3],
+                   pretrained, progress, **kwargs)
+
+
+def resnext101_32x8d(pretrained=False, progress=True, **kwargs):
+    r"""ResNeXt-101 32x8d model from
+    `"Aggregated Residual Transformation for Deep Neural Networks" <https://arxiv.org/pdf/1611.05431.pdf>`_
+
+    Args:
+        pretrained (bool): If True, returns a model pre-trained on ImageNet
+        progress (bool): If True, displays a progress bar of the download to stderr
+    """
+    kwargs['groups'] = 32
+    kwargs['width_per_group'] = 8
+    return _resnet('resnext101_32x8d', Bottleneck, [3, 4, 23, 3],
+                   pretrained, progress, **kwargs)
+
+
+def wide_resnet50_2(pretrained=False, progress=True, **kwargs):
+    r"""Wide ResNet-50-2 model from
+    `"Wide Residual Networks" <https://arxiv.org/pdf/1605.07146.pdf>`_
+
+    The model is the same as ResNet except for the bottleneck number of channels
+    which is twice larger in every block. The number of channels in outer 1x1
+    convolutions is the same, e.g. last block in ResNet-50 has 2048-512-2048
+    channels, and in Wide ResNet-50-2 has 2048-1024-2048.
+
+    Args:
+        pretrained (bool): If True, returns a model pre-trained on ImageNet
+        progress (bool): If True, displays a progress bar of the download to stderr
+    """
+    kwargs['width_per_group'] = 64 * 2
+    return _resnet('wide_resnet50_2', Bottleneck, [3, 4, 6, 3],
+                   pretrained, progress, **kwargs)
+
+
+def wide_resnet50_3(pretrained=False, progress=True, **kwargs):
+    r"""Wide ResNet-50-3 model
+    Args:
+        pretrained (bool): If True, returns a model pre-trained on ImageNet
+        progress (bool): If True, displays a progress bar of the download to stderr
+    """
+    kwargs['width_per_group'] = 64 * 3
+    return _resnet('wide_resnet50_3', Bottleneck, [3, 4, 6, 3],
+                   pretrained, progress, **kwargs)
+
+
+def wide_resnet50_4(pretrained=False, progress=True, **kwargs):
+    r"""Wide ResNet-50-4 model
+    Args:
+        pretrained (bool): If True, returns a model pre-trained on ImageNet
+        progress (bool): If True, displays a progress bar of the download to stderr
+    """
+    kwargs['width_per_group'] = 64 * 4
+    return _resnet('wide_resnet50_4', Bottleneck, [3, 4, 6, 3],
+                   pretrained, progress, **kwargs)
+
+
+def wide_resnet50_5(pretrained=False, progress=True, **kwargs):
+    r"""Wide ResNet-50-5 model
+    Args:
+        pretrained (bool): If True, returns a model pre-trained on ImageNet
+        progress (bool): If True, displays a progress bar of the download to stderr
+    """
+    kwargs['width_per_group'] = 64 * 5
+    return _resnet('wide_resnet50_5', Bottleneck, [3, 4, 6, 3],
+                   pretrained, progress, **kwargs)
+
+
+def wide_resnet50_6(pretrained=False, progress=True, **kwargs):
+    r"""Wide ResNet-50-6 model
+    Args:
+        pretrained (bool): If True, returns a model pre-trained on ImageNet
+        progress (bool): If True, displays a progress bar of the download to stderr
+    """
+    kwargs['width_per_group'] = 64 * 6
+    return _resnet('wide_resnet50_6', Bottleneck, [3, 4, 6, 3],
+                   pretrained, progress, **kwargs)
+
+
+def wide_resnet101_2(pretrained=False, progress=True, **kwargs):
+    r"""Wide ResNet-101-2 model from
+    `"Wide Residual Networks" <https://arxiv.org/pdf/1605.07146.pdf>`_
+
+    The model is the same as ResNet except for the bottleneck number of channels
+    which is twice larger in every block. The number of channels in outer 1x1
+    convolutions is the same, e.g. last block in ResNet-50 has 2048-512-2048
+    channels, and in Wide ResNet-50-2 has 2048-1024-2048.
+
+    Args:
+        pretrained (bool): If True, returns a model pre-trained on ImageNet
+        progress (bool): If True, displays a progress bar of the download to stderr
+    """
+    kwargs['width_per_group'] = 64 * 2
+    return _resnet('wide_resnet101_2', Bottleneck, [3, 4, 23, 3],
+                   pretrained, progress, **kwargs)

architectures/utils.py

@@ -0,0 +1,17 @@
+import torch
+
+
+
+class SequentialWithArgs(torch.nn.Sequential):
+    def forward(self, input, *args, **kwargs):
+        vs = list(self._modules.values())
+        l = len(vs)
+        for i in range(l):
+            if i == l-1:
+                input = vs[i](input, *args, **kwargs)
+            else:
+                input = vs[i](input)
+        return input
+
+
+

configs/architecture_params.py

@@ -0,0 +1 @@
+architecture_params = {"resnet50": {"beta":0.196}}

configs/dataset_params.py

@@ -0,0 +1,22 @@
+import torch
+
+from configs.optim_params import EvaluatedDict
+
+dataset_constants = {"CUB2011":{"num_classes":200},
+                      "TravelingBirds":{"num_classes":200},
+                      "ImageNet":{"num_classes":1000},
+                      "StanfordCars":{"num_classes":196},
+                     "FGVCAircraft": {"num_classes":100}}
+
+normalize_params = {"CUB2011":{"mean":  torch.tensor([0.4853, 0.4964, 0.4295]),"std":torch.tensor([0.2300, 0.2258, 0.2625])},
+"TravelingBirds":{"mean":  torch.tensor([0.4584, 0.4369, 0.3957]),"std":torch.tensor([0.2610, 0.2569, 0.2722])},
+                    "ImageNet":{'mean': torch.tensor([0.485, 0.456, 0.406]),'std': torch.tensor([0.229, 0.224, 0.225])} ,
+"StanfordCars":{'mean': torch.tensor([0.4593, 0.4466, 0.4453]),'std': torch.tensor([0.2920, 0.2910, 0.2988])} ,
+                  "FGVCAircraft":{'mean': torch.tensor([0.4827, 0.5130, 0.5352]),
+                    'std': torch.tensor([0.2236, 0.2170, 0.2478]),}
+                    }
+
+
+dense_batch_size =  EvaluatedDict({False: 16,True: 1024,}, lambda x: x == "ImageNet")
+
+ft_batch_size =  EvaluatedDict({False: 16,True: 1024,}, lambda x: x == "ImageNet")# Untested

configs/optim_params.py

@@ -0,0 +1,22 @@
+# order: lr,weight_decay, step_lr, step_lr_gamma
+import math
+
+
+class EvaluatedDict:
+    def __init__(self, d, func):
+        self.dict = d
+        self.func = func
+
+    def __getitem__(self, key):
+        return self.dict[self.func(key)]
+
+dense_params = EvaluatedDict({False: [0.005, 0.0005, 30, 0.4, 150],True: [None,None,None,None,None],}, lambda x: x == "ImageNet")
+def calculate_lr_from_args( epochs, step_lr, start_lr, step_lr_decay):
+    # Gets the final learning rate after dense training with step_lr_schedule.
+    n_steps = math.floor((epochs - step_lr) / step_lr)
+    final_lr = start_lr * step_lr_decay ** n_steps
+    return final_lr
+
+ft_params =EvaluatedDict({False: [1e-4, 0.0005, 10, 0.4, 40],True:[[calculate_lr_from_args(150,30,0.005, 0.4), 0.0005, 10, 0.4, 40]]}, lambda x: x == "ImageNet")
+
+

configs/qsenn_training_params.py

@@ -0,0 +1,11 @@
+from configs.sldd_training_params import OptimizationScheduler
+
+
+class  QSENNScheduler(OptimizationScheduler):
+    def get_params(self):
+        params = super().get_params()
+        if self.n_calls >= 2:
+            params[0] = params[0] * 0.9**(self.n_calls-2)
+        if 2 <= self.n_calls <= 4:
+            params[-2] = 10# Change num epochs to 10 for iterative finetuning
+        return params

configs/sldd_training_params.py

@@ -0,0 +1,17 @@
+from configs.optim_params import dense_params, ft_params
+
+
+class OptimizationScheduler:
+    def __init__(self, dataset):
+        self.dataset = dataset
+        self.n_calls =  0
+
+
+    def get_params(self):
+        if self.n_calls == 0: # Return Deńse Params
+            params =  dense_params[self.dataset]+ [False]
+        else: # Return Finetuning Params
+            params =  ft_params[self.dataset]+ [True]
+        self.n_calls += 1
+        return params
+

dataset_classes/cub200.py

@@ -0,0 +1,96 @@
+# Dataset should lie under /root/
+# root is currently set to ~/tmp/Datasets/CUB200
+# If cropped iamges, like for PIP-Net, ProtoPool, etc. are used, then the crop_root should be set to a folder containing the
+# cropped images in the expected structure, obtained by following ProtoTree's instructions.
+# https://github.com/M-Nauta/ProtoTree/blob/main/README.md#preprocessing-cub
+import os
+from pathlib import Path
+
+import numpy as np
+import pandas as pd
+from torch.utils.data import Dataset
+from torchvision.datasets.folder import default_loader
+
+from dataset_classes.utils import txt_load
+
+
+class CUB200Class(Dataset):
+    root = Path.home() / "tmp/Datasets/CUB200"
+    crop_root = Path.home() / "tmp/Datasets/PPCUB200"
+    base_folder = 'CUB_200_2011/images'
+    def __init__(self,  train, transform, crop=True):
+        self.train = train
+        self.transform = transform
+        self.crop = crop
+        self._load_metadata()
+        self.loader = default_loader
+
+        if crop:
+            self.adapt_to_crop()
+
+    def _load_metadata(self):
+        images = pd.read_csv(os.path.join(self.root, 'CUB_200_2011', 'images.txt'), sep=' ',
+                             names=['img_id', 'filepath'])
+        image_class_labels = pd.read_csv(os.path.join(self.root, 'CUB_200_2011', 'image_class_labels.txt'),
+                                         sep=' ', names=['img_id', 'target'])
+        train_test_split = pd.read_csv(os.path.join(self.root, 'CUB_200_2011', 'train_test_split.txt'),
+                                       sep=' ', names=['img_id', 'is_training_img'])
+        data = images.merge(image_class_labels, on='img_id')
+        self.data = data.merge(train_test_split, on='img_id')
+        if self.train:
+            self.data = self.data[self.data.is_training_img == 1]
+        else:
+            self.data = self.data[self.data.is_training_img == 0]
+
+    def __len__(self):
+        return len(self.data)
+
+    def adapt_to_crop(self):
+      #  ds_name = [x for x in self.cropped_dict.keys() if x in self.root][0]
+        self.root = self.crop_root
+        folder_name = "train" if self.train else "test"
+        folder_name = folder_name + "_cropped"
+        self.base_folder = 'CUB_200_2011/' + folder_name
+
+    def __getitem__(self, idx):
+        sample = self.data.iloc[idx]
+        path = os.path.join(self.root, self.base_folder, sample.filepath)
+        target = sample.target - 1  # Targets start at 1 by default, so shift to 0
+        img = self.loader(path)
+        img = self.transform(img)
+        return img, target
+
+    @classmethod
+    def get_image_attribute_labels(self, train=False):
+        image_attribute_labels = pd.read_csv(
+            os.path.join('/home/norrenbr/tmp/Datasets/CUB200', 'CUB_200_2011', "attributes",
+                         'image_attribute_labels.txt'),
+            sep=' ', names=['img_id', 'attribute', "is_present", "certainty", "time"], on_bad_lines="skip")
+        train_test_split = pd.read_csv(os.path.join(self.root, 'CUB_200_2011', 'train_test_split.txt'),
+                                       sep=' ', names=['img_id', 'is_training_img'])
+        merged = image_attribute_labels.merge(train_test_split, on="img_id")
+        filtered_data = merged[merged["is_training_img"] == train]
+        return filtered_data
+
+
+    @staticmethod
+    def filter_attribute_labels(labels, min_certainty=3):
+        is_invisible_present = labels[labels["certainty"] == 1]["is_present"].sum()
+        if is_invisible_present != 0:
+            raise ValueError("Invisible present")
+        labels["img_id"] -= min(labels["img_id"])
+        labels["img_id"] = fillholes_in_array(labels["img_id"])
+        labels[labels["certainty"] == 1]["certainty"] = 4
+        labels = labels[labels["certainty"] >= min_certainty]
+        labels["attribute"] -= min(labels["attribute"])
+        labels = labels[["img_id", "attribute", "is_present"]]
+        labels["is_present"] = labels["is_present"].astype(bool)
+        return labels
+
+
+
+def fillholes_in_array(array):
+    unique_values = np.unique(array)
+    mapping = {x: i for i, x in enumerate(unique_values)}
+    array = array.map(mapping)
+    return array

dataset_classes/stanfordcars.py

@@ -0,0 +1,121 @@
+import pathlib
+from typing import Callable, Optional, Any, Tuple
+
+import numpy as np
+import pandas as pd
+from PIL import Image
+from torchvision.datasets import VisionDataset
+from torchvision.datasets.utils import download_and_extract_archive, download_url
+
+
+class StanfordCarsClass(VisionDataset):
+    """`Stanford Cars <https://ai.stanford.edu/~jkrause/cars/car_dataset.html>`_ Dataset
+
+    The Cars dataset contains 16,185 images of 196 classes of cars. The data is
+    split into 8,144 training images and 8,041 testing images, where each class
+    has been split roughly in a 50-50 split
+
+    .. note::
+
+        This class needs `scipy <https://docs.scipy.org/doc/>`_ to load target files from `.mat` format.
+
+    Args:
+        root (string): Root directory of dataset
+        split (string, optional): The dataset split, supports ``"train"`` (default) or ``"test"``.
+        transform (callable, optional): A function/transform that  takes in an PIL image
+            and returns a transformed version. E.g, ``transforms.RandomCrop``
+        target_transform (callable, optional): A function/transform that takes in the
+            target and transforms it.
+        download (bool, optional): If True, downloads the dataset from the internet and
+            puts it in root directory. If dataset is already downloaded, it is not
+            downloaded again."""
+    root = pathlib.Path.home() / "tmp" / "Datasets" / "StanfordCars"
+    def __init__(
+            self,
+            train: bool = True,
+            transform: Optional[Callable] = None,
+            target_transform: Optional[Callable] = None,
+            download: bool = True,
+    ) -> None:
+
+        try:
+            import scipy.io as sio
+        except ImportError:
+            raise RuntimeError("Scipy is not found. This dataset needs to have scipy installed: pip install scipy")
+
+        super().__init__(self.root, transform=transform, target_transform=target_transform)
+
+        self.train = train
+        self._base_folder = pathlib.Path(self.root) / "stanford_cars"
+        devkit = self._base_folder / "devkit"
+
+        if train:
+            self._annotations_mat_path = devkit / "cars_train_annos.mat"
+            self._images_base_path = self._base_folder / "cars_train"
+        else:
+            self._annotations_mat_path = self._base_folder / "cars_test_annos_withlabels.mat"
+            self._images_base_path = self._base_folder / "cars_test"
+
+        if download:
+            self.download()
+
+        if not self._check_exists():
+            raise RuntimeError("Dataset not found. You can use download=True to download it")
+
+        self.samples = [
+            (
+                str(self._images_base_path / annotation["fname"]),
+                annotation["class"] - 1,  # Original target mapping  starts from 1, hence -1
+            )
+            for annotation in sio.loadmat(self._annotations_mat_path, squeeze_me=True)["annotations"]
+        ]
+        self.targets = np.array([x[1] for x in self.samples])
+        self.classes = sio.loadmat(str(devkit / "cars_meta.mat"), squeeze_me=True)["class_names"].tolist()
+        self.class_to_idx = {cls: i for i, cls in enumerate(self.classes)}
+
+    def __len__(self) -> int:
+        return len(self.samples)
+
+    def __getitem__(self, idx: int) -> Tuple[Any, Any]:
+        """Returns pil_image and class_id for given index"""
+        image_path, target = self.samples[idx]
+        pil_image = Image.open(image_path).convert("RGB")
+
+        if self.transform is not None:
+            pil_image = self.transform(pil_image)
+        if self.target_transform is not None:
+            target = self.target_transform(target)
+        return pil_image, target
+
+    def download(self) -> None:
+        if self._check_exists():
+            return
+
+        download_and_extract_archive(
+            url="https://ai.stanford.edu/~jkrause/cars/car_devkit.tgz",
+            download_root=str(self._base_folder),
+            md5="c3b158d763b6e2245038c8ad08e45376",
+        )
+        if self.train:
+            download_and_extract_archive(
+                url="https://ai.stanford.edu/~jkrause/car196/cars_train.tgz",
+                download_root=str(self._base_folder),
+                md5="065e5b463ae28d29e77c1b4b166cfe61",
+            )
+        else:
+            download_and_extract_archive(
+                url="https://ai.stanford.edu/~jkrause/car196/cars_test.tgz",
+                download_root=str(self._base_folder),
+                md5="4ce7ebf6a94d07f1952d94dd34c4d501",
+            )
+            download_url(
+                url="https://ai.stanford.edu/~jkrause/car196/cars_test_annos_withlabels.mat",
+                root=str(self._base_folder),
+                md5="b0a2b23655a3edd16d84508592a98d10",
+            )
+
+    def _check_exists(self) -> bool:
+        if not (self._base_folder / "devkit").is_dir():
+            return False
+
+        return self._annotations_mat_path.exists() and self._images_base_path.is_dir()

dataset_classes/travelingbirds.py

@@ -0,0 +1,59 @@
+# TravelingBirds dataset needs to be downloaded from https://worksheets.codalab.org/bundles/0x518829de2aa440c79cd9d75ef6669f27
+# as it comes from https://github.com/yewsiang/ConceptBottleneck
+import os
+from pathlib import Path
+
+import numpy as np
+import pandas as pd
+
+from dataset_classes.cub200 import CUB200Class
+from dataset_classes.utils import index_list_with_sorting, mask_list
+
+
+class TravelingBirds(CUB200Class):
+    init_base_folder = 'CUB_fixed'
+    root = Path.home() / "tmp/Datasets/TravelingBirds"
+    crop_root = Path.home() / "tmp/Datasets/PPTravelingBirds"
+    def get_all_samples_dir(self, dir):
+
+        self.base_folder = os.path.join(self.init_base_folder, dir)
+        main_dir = Path(self.root) / self.init_base_folder / dir
+        return self.get_all_sample(main_dir)
+
+    def adapt_to_crop(self):
+        self.root = self.crop_root
+        folder_name = "train" if self.train else "test"
+        folder_name = folder_name + "_cropped"
+        self.base_folder = 'CUB_fixed/' + folder_name
+
+    def get_all_sample(self, dir):
+        answer = []
+        for i, sub_dir in enumerate(sorted(os.listdir(dir))):
+            class_dir = dir / sub_dir
+            for single_img in os.listdir(class_dir):
+                answer.append([Path(sub_dir) / single_img, i + 1])
+        return answer
+    def _load_metadata(self):
+        train_test_split = pd.read_csv(
+            os.path.join(Path(self.root).parent / "CUB200", 'CUB_200_2011', 'train_test_split.txt'),
+            sep=' ', names=['img_id', 'is_training_img'])
+        data = pd.read_csv(
+            os.path.join(Path(self.root).parent / "CUB200", 'CUB_200_2011', 'images.txt'),
+            sep=' ', names=['img_id', "path"])
+        img_dict = {x[1]: x[0] for x in data.values}
+        # TravelingBirds has all train+test images in both folders, just with different backgrounds.
+        # They are separated by train_test_split of CUB200.
+        if self.train:
+            samples = self.get_all_samples_dir("train")
+            mask = train_test_split["is_training_img"] == 1
+        else:
+            samples = self.get_all_samples_dir("test")
+            mask = train_test_split["is_training_img"] == 0
+        ids = np.array([img_dict[str(x[0])] for x in samples])
+        sorted = np.argsort(ids)
+        samples = index_list_with_sorting(samples, sorted)
+        samples = mask_list(samples, mask)
+        filepaths = [x[0] for x in samples]
+        labels = [x[1] for x in samples]
+        samples = pd.DataFrame({"filepath": filepaths, "target": labels})
+        self.data = samples

dataset_classes/utils.py

@@ -0,0 +1,16 @@
+def index_list_with_sorting(list_to_sort, sorting_list):
+    answer = []
+    for entry in sorting_list:
+        answer.append(list_to_sort[entry])
+    return answer
+
+
+def mask_list(list_input, mask):
+    return [x for i, x in enumerate(list_input) if mask[i]]
+
+
+def txt_load(filename):
+    with open(filename, 'r') as f:
+        data = f.read()
+    return data
+

environment.yml

@@ -0,0 +1,117 @@
+name: QSENNEnv
+channels:
+  - pytorch
+  - nvidia
+  - defaults
+dependencies:
+  - _libgcc_mutex=0.1=main
+  - _openmp_mutex=5.1=1_gnu
+  - blas=1.0=mkl
+  - brotli-python=1.0.9=py310h6a678d5_7
+  - bzip2=1.0.8=h7b6447c_0
+  - ca-certificates=2023.12.12=h06a4308_0
+  - certifi=2023.11.17=py310h06a4308_0
+  - cffi=1.16.0=py310h5eee18b_0
+  - charset-normalizer=2.0.4=pyhd3eb1b0_0
+  - cryptography=41.0.7=py310hdda0065_0
+  - cuda-cudart=12.1.105=0
+  - cuda-cupti=12.1.105=0
+  - cuda-libraries=12.1.0=0
+  - cuda-nvrtc=12.1.105=0
+  - cuda-nvtx=12.1.105=0
+  - cuda-opencl=12.3.101=0
+  - cuda-runtime=12.1.0=0
+  - ffmpeg=4.3=hf484d3e_0
+  - filelock=3.13.1=py310h06a4308_0
+  - freetype=2.12.1=h4a9f257_0
+  - giflib=5.2.1=h5eee18b_3
+  - gmp=6.2.1=h295c915_3
+  - gmpy2=2.1.2=py310heeb90bb_0
+  - gnutls=3.6.15=he1e5248_0
+  - idna=3.4=py310h06a4308_0
+  - intel-openmp=2023.1.0=hdb19cb5_46306
+  - jinja2=3.1.2=py310h06a4308_0
+  - jpeg=9e=h5eee18b_1
+  - lame=3.100=h7b6447c_0
+  - lcms2=2.12=h3be6417_0
+  - ld_impl_linux-64=2.38=h1181459_1
+  - lerc=3.0=h295c915_0
+  - libcublas=12.1.0.26=0
+  - libcufft=11.0.2.4=0
+  - libcufile=1.8.1.2=0
+  - libcurand=10.3.4.107=0
+  - libcusolver=11.4.4.55=0
+  - libcusparse=12.0.2.55=0
+  - libdeflate=1.17=h5eee18b_1
+  - libffi=3.4.4=h6a678d5_0
+  - libgcc-ng=11.2.0=h1234567_1
+  - libgomp=11.2.0=h1234567_1
+  - libiconv=1.16=h7f8727e_2
+  - libidn2=2.3.4=h5eee18b_0
+  - libjpeg-turbo=2.0.0=h9bf148f_0
+  - libnpp=12.0.2.50=0
+  - libnvjitlink=12.1.105=0
+  - libnvjpeg=12.1.1.14=0
+  - libpng=1.6.39=h5eee18b_0
+  - libstdcxx-ng=11.2.0=h1234567_1
+  - libtasn1=4.19.0=h5eee18b_0
+  - libtiff=4.5.1=h6a678d5_0
+  - libunistring=0.9.10=h27cfd23_0
+  - libuuid=1.41.5=h5eee18b_0
+  - libwebp=1.3.2=h11a3e52_0
+  - libwebp-base=1.3.2=h5eee18b_0
+  - llvm-openmp=14.0.6=h9e868ea_0
+  - lz4-c=1.9.4=h6a678d5_0
+  - markupsafe=2.1.3=py310h5eee18b_0
+  - mkl=2023.1.0=h213fc3f_46344
+  - mkl-service=2.4.0=py310h5eee18b_1
+  - mkl_fft=1.3.8=py310h5eee18b_0
+  - mkl_random=1.2.4=py310hdb19cb5_0
+  - mpc=1.1.0=h10f8cd9_1
+  - mpfr=4.0.2=hb69a4c5_1
+  - mpmath=1.3.0=py310h06a4308_0
+  - ncurses=6.4=h6a678d5_0
+  - nettle=3.7.3=hbbd107a_1
+  - networkx=3.1=py310h06a4308_0
+  - numpy=1.26.3=py310h5f9d8c6_0
+  - numpy-base=1.26.3=py310hb5e798b_0
+  - openh264=2.1.1=h4ff587b_0
+  - openjpeg=2.4.0=h3ad879b_0
+  - openssl=3.0.12=h7f8727e_0
+  - pillow=10.0.1=py310ha6cbd5a_0
+  - pip=23.3.1=py310h06a4308_0
+  - pycparser=2.21=pyhd3eb1b0_0
+  - pyopenssl=23.2.0=py310h06a4308_0
+  - pysocks=1.7.1=py310h06a4308_0
+  - python=3.10.13=h955ad1f_0
+  - pytorch=2.1.2=py3.10_cuda12.1_cudnn8.9.2_0
+  - pytorch-cuda=12.1=ha16c6d3_5
+  - pytorch-mutex=1.0=cuda
+  - pyyaml=6.0.1=py310h5eee18b_0
+  - readline=8.2=h5eee18b_0
+  - requests=2.31.0=py310h06a4308_0
+  - setuptools=68.2.2=py310h06a4308_0
+  - sqlite=3.41.2=h5eee18b_0
+  - sympy=1.12=py310h06a4308_0
+  - tbb=2021.8.0=hdb19cb5_0
+  - tk=8.6.12=h1ccaba5_0
+  - torchaudio=2.1.2=py310_cu121
+  - torchtriton=2.1.0=py310
+  - torchvision=0.16.2=py310_cu121
+  - typing_extensions=4.7.1=py310h06a4308_0
+  - urllib3=1.26.18=py310h06a4308_0
+  - wheel=0.41.2=py310h06a4308_0
+  - xz=5.4.5=h5eee18b_0
+  - yaml=0.2.5=h7b6447c_0
+  - zlib=1.2.13=h5eee18b_0
+  - zstd=1.5.5=hc292b87_0
+  - pip:
+      - fsspec==2023.12.2
+      - glm-saga==0.1.2
+      - pandas==2.1.4
+      - python-dateutil==2.8.2
+      - pytz==2023.3.post1
+      - six==1.16.0
+      - tqdm==4.66.1
+      - tzdata==2023.4
+prefix: /home/norrenbr/anaconda/tmp/envs/QSENN-Minimal

evaluation/Metrics/Dependence.py

@@ -0,0 +1,21 @@
+import torch
+
+
+def compute_contribution_top_feature(features, outputs, weights,  labels):
+    with torch.no_grad():
+        total_pre_softmax, predicted_classes = torch.max(outputs, dim=1)
+        feature_part = features * weights.to(features.device)[predicted_classes]
+        class_specific_feature_part = torch.zeros((weights.shape[0], features.shape[1],))
+        feature_class_part = torch.zeros((weights.shape[0], features.shape[1],))
+        for unique_class in predicted_classes.unique():
+            mask = predicted_classes == unique_class
+            class_specific_feature_part[unique_class] = feature_part[mask].mean(dim=0)
+            gt_mask = labels == unique_class
+            feature_class_part[unique_class] = feature_part[gt_mask].mean(dim=0)
+        abs_features = feature_part.abs()
+        abs_sum = abs_features.sum(dim=1)
+        fractions_abs = abs_features / abs_sum[:, None]
+        abs_max = fractions_abs.max(dim=1)[0]
+        mask = ~torch.isnan(abs_max)
+        abs_max = abs_max[mask]
+    return abs_max.mean()

evaluation/Metrics/cub_Alignment.py

@@ -0,0 +1,30 @@
+import numpy as np
+
+from dataset_classes.cub200 import CUB200Class
+
+
+def get_cub_alignment_from_features(features_train_sorted):
+    metric_matrix = compute_metric_matrix(np.array(features_train_sorted), "train")
+    return np.mean(np.max(metric_matrix, axis=1))
+    pass
+
+
+def compute_metric_matrix(features, mode):
+    image_attribute_labels = CUB200Class.get_image_attribute_labels(train=mode == "train")
+    image_attribute_labels = CUB200Class.filter_attribute_labels(image_attribute_labels)
+    matrix_shape = (
+        features.shape[1], max(image_attribute_labels["attribute"]) + 1)
+    accuracy_matrix = np.zeros(matrix_shape)
+    sensitivity_matrix = np.zeros_like(accuracy_matrix)
+    grouped_attributes = image_attribute_labels.groupby("attribute")
+    for attribute_id, group in grouped_attributes:
+        is_present = group[group["is_present"]]
+        not_present = group[~group["is_present"]]
+        is_present_avg = np.mean(features[is_present["img_id"]], axis=0)
+        not_present_avg = np.mean(features[not_present["img_id"]], axis=0)
+        sensitivity_matrix[:, attribute_id] = not_present_avg
+        accuracy_matrix[:, attribute_id] = is_present_avg
+    metric_matrix = accuracy_matrix - sensitivity_matrix
+    no_abs_features = features - np.min(features, axis=0)
+    no_abs_feature_mean = metric_matrix / no_abs_features.mean(axis=0)[:, None]
+    return  no_abs_feature_mean

evaluation/diversity.py

@@ -0,0 +1,111 @@
+import numpy as np
+import torch
+
+from evaluation.helpers import softmax_feature_maps
+
+
+class MultiKCrossChannelMaxPooledSum:
+    def __init__(self, top_k_range, weights, interactions, func="softmax"):
+        self.top_k_range = top_k_range
+        self.weights = weights
+        self.failed = False
+        self.max_ks = self.get_max_ks(weights)
+        self.locality_of_used_features = torch.zeros(len(top_k_range), device=weights.device)
+        self.locality_of_exclusely_used_features = torch.zeros(len(top_k_range), device=weights.device)
+        self.ns_k = torch.zeros(len(top_k_range), device=weights.device)
+        self.exclusive_ns = torch.zeros(len(top_k_range), device=weights.device)
+        self.interactions = interactions
+        self.func = func
+
+    def get_max_ks(self, weights):
+        nonzeros = torch.count_nonzero(torch.tensor(weights), 1)
+        return nonzeros
+
+    def get_top_n_locality(self, outputs, initial_feature_maps, k):
+        feature_maps, relevant_weights, vector_size, top_classes = self.adapt_feature_maps(outputs,
+                                                                                           initial_feature_maps)
+        max_ks = self.max_ks[top_classes]
+        max_k_based_row_selection = max_ks >= k
+
+        result = self.get_crosspooled(relevant_weights, max_k_based_row_selection, k, vector_size, feature_maps,
+                                      separated=True)
+        return result
+
+    def get_locality(self, outputs, initial_feature_maps, n):
+        answer = self.get_top_n_locality(outputs, initial_feature_maps, n)
+        return answer
+
+    def get_result(self):
+        # if torch.sum(self.exclusive_ns) ==0:
+        #     end_idx = len(self.exclusive_ns) - 1
+        # else:
+
+        exclusive_array = torch.zeros_like(self.locality_of_exclusely_used_features)
+        local_array = torch.zeros_like(self.locality_of_used_features)
+        # if self.failed:
+        #     return local_array, exclusive_array
+        cumulated = torch.cumsum(self.exclusive_ns, 0)
+        end_idx = torch.argmax(cumulated)
+        exclusivity_array = self.locality_of_exclusely_used_features[:end_idx + 1] / self.exclusive_ns[:end_idx + 1]
+        exclusivity_array[exclusivity_array != exclusivity_array] = 0
+        exclusive_array[:len(exclusivity_array)] = exclusivity_array
+        locality_array = self.locality_of_used_features[self.locality_of_used_features != 0] / self.ns_k[
+            self.locality_of_used_features != 0]
+        local_array[:len(locality_array)] = locality_array
+        return local_array, exclusive_array
+
+    def get_crosspooled(self, relevant_weights, mask, k, vector_size, feature_maps, separated=False):
+        relevant_indices = get_relevant_indices(relevant_weights, k)[mask]
+        # this should have size batch x k x featuremapsize squared]
+        indices = relevant_indices.unsqueeze(2).repeat(1, 1, vector_size)
+        sub_feature_maps = torch.gather(feature_maps[mask], 1, indices)
+        # shape batch x featuremapsquared: For each "pixel" the highest value
+        cross_pooled = torch.max(sub_feature_maps, 1)[0]
+        if separated:
+            return torch.sum(cross_pooled, 1) / k
+        else:
+            ns = len(cross_pooled)
+            result = torch.sum(cross_pooled) / (k)
+            # should be batch x map size
+
+            return ns, result
+
+    def adapt_feature_maps(self, outputs, initial_feature_maps):
+        if self.func == "softmax":
+            feature_maps = softmax_feature_maps(initial_feature_maps)
+        feature_maps = torch.flatten(feature_maps, 2)
+        vector_size = feature_maps.shape[2]
+        top_classes = torch.argmax(outputs, dim=1)
+        relevant_weights = self.weights[top_classes]
+        if relevant_weights.shape[1] != feature_maps.shape[1]:
+            feature_maps = self.interactions.get_localized_features(initial_feature_maps)
+            feature_maps = softmax_feature_maps(feature_maps)
+            feature_maps = torch.flatten(feature_maps, 2)
+        return feature_maps, relevant_weights, vector_size, top_classes
+
+    def calculate_locality(self, outputs, initial_feature_maps):
+        feature_maps, relevant_weights, vector_size, top_classes = self.adapt_feature_maps(outputs,
+                                                                                           initial_feature_maps)
+        max_ks = self.max_ks[top_classes]
+        for k in self.top_k_range:
+            # relevant_k_s = max_ks[]
+            max_k_based_row_selection = max_ks >= k
+            if torch.sum(max_k_based_row_selection) == 0:
+                break
+
+            exclusive_k = max_ks == k
+            if torch.sum(exclusive_k) != 0:
+                ns, result = self.get_crosspooled(relevant_weights, exclusive_k, k, vector_size, feature_maps)
+                self.locality_of_exclusely_used_features[k - 1] += result
+                self.exclusive_ns[k - 1] += ns
+            ns, result = self.get_crosspooled(relevant_weights, max_k_based_row_selection, k, vector_size, feature_maps)
+            self.ns_k[k - 1] += ns
+            self.locality_of_used_features[k - 1] += result
+
+    def __call__(self, outputs, initial_feature_maps):
+        self.calculate_locality(outputs, initial_feature_maps)
+
+
+def get_relevant_indices(weights, top_k):
+    top_k = weights.topk(top_k)[1]
+    return top_k

evaluation/helpers.py

@@ -0,0 +1,6 @@
+import torch
+
+
+def softmax_feature_maps(x):
+    # done: verify that this applies softmax along first dimension
+    return torch.softmax(x.reshape(x.size(0), x.size(1), -1), 2).view_as(x)

evaluation/qsenn_metrics.py

@@ -0,0 +1,39 @@
+import numpy as np
+import torch
+
+from evaluation.Metrics.Dependence import compute_contribution_top_feature
+from evaluation.Metrics.cub_Alignment import get_cub_alignment_from_features
+from evaluation.diversity import MultiKCrossChannelMaxPooledSum
+from evaluation.utils import get_metrics_for_model
+
+
+def evaluateALLMetricsForComps(features_train,  outputs_train,  feature_maps_test,
+                               outputs_test, linear_matrix,  labels_train):
+    with torch.no_grad():
+        if len(features_train) < 7000: # recognize CUB and TravelingBirds
+            cub_alignment = get_cub_alignment_from_features(features_train)
+        else:
+            cub_alignment = 0
+        print("cub_alignment: ", cub_alignment)
+        localizer = MultiKCrossChannelMaxPooledSum(range(1, 6), linear_matrix, None)
+        batch_size = 300
+        for i in range(np.floor(len(features_train) / batch_size).astype(int)):
+            localizer(outputs_test[i * batch_size:(i + 1) * batch_size].to("cuda"),
+                      feature_maps_test[i * batch_size:(i + 1) * batch_size].to("cuda"))
+
+        locality, exlusive_locality = localizer.get_result()
+        diversity = locality[4]
+        print("diversity@5: ", diversity)
+        abs_frac_mean = compute_contribution_top_feature(
+            features_train,
+            outputs_train,
+            linear_matrix,
+     labels_train)
+        print("Dependence ", abs_frac_mean)
+        answer_dict = {"diversity": diversity.item(),  "Dependence": abs_frac_mean.item(), "Alignment":cub_alignment}
+    return answer_dict
+
+def eval_model_on_all_qsenn_metrics(model, test_loader, train_loader):
+    return get_metrics_for_model(train_loader, test_loader, model, evaluateALLMetricsForComps)
+
+

evaluation/utils.py

@@ -0,0 +1,57 @@
+import torch
+from tqdm import tqdm
+
+
+
+def get_metrics_for_model(train_loader, test_loader, model, metric_evaluator):
+    (features_train, feature_maps_train, outputs_train, features_test, feature_maps_test,
+     outputs_test, labels) = [], [], [], [], [], [], []
+    device = "cuda" if torch.cuda.is_available() else "cpu"
+    model.eval()
+    model = model.to(device)
+    training_transforms = train_loader.dataset.transform
+    train_loader.dataset.transform = test_loader.dataset.transform # Use test transform for train
+    train_loader = torch.utils.data.DataLoader(train_loader.dataset, batch_size=100, shuffle=False) # Turn off shuffling
+    print("Going in get metrics")
+    linear_matrix = model.linear.weight
+    entries = torch.nonzero(linear_matrix)
+    rel_features = torch.unique(entries[:, 1])
+    with torch.no_grad():
+        iterator = tqdm(enumerate(train_loader), total=len(train_loader))
+        for batch_idx, (data, target) in iterator:
+            xs1 = data.to("cuda")
+            output, feature_maps,  final_features = model(xs1,   with_feature_maps=True,                                                                                                              with_final_features=True,)
+            outputs_train.append(output.to("cpu"))
+            features_train.append(final_features.to("cpu"))
+            labels.append(target.to("cpu"))
+        total = 0
+        correct = 0
+        iterator = tqdm(enumerate(test_loader), total=len(test_loader))
+        for batch_idx, (data, target) in iterator:
+            xs1 = data.to("cuda")
+            output, feature_maps, final_features = model(xs1, with_feature_maps=True,
+                                                         with_final_features=True, )
+            feature_maps_test.append(feature_maps[:, rel_features].to("cpu"))
+            outputs_test.append(output.to("cpu"))
+            total += target.size(0)
+            _, predicted = output.max(1)
+            correct += predicted.eq(target.to("cuda")).sum().item()
+        print("test accuracy: ", correct / total)
+        features_train = torch.cat(features_train)
+        outputs_train = torch.cat(outputs_train)
+        feature_maps_test = torch.cat(feature_maps_test)
+        outputs_test = torch.cat(outputs_test)
+        labels = torch.cat(labels)
+        linear_matrix = linear_matrix[:, rel_features]
+    print("Shape of linear matrix: ", linear_matrix.shape)
+    all_metrics_dict = metric_evaluator(features_train,  outputs_train,
+                                             feature_maps_test,
+                                             outputs_test, linear_matrix,  labels)
+    result_dict = {"Accuracy": correct / total,  "NFfeatures": linear_matrix.shape[1],
+                   "PerClass": torch.nonzero(linear_matrix).shape[0] / linear_matrix.shape[0],
+                   }
+    result_dict.update(all_metrics_dict)
+    print(result_dict)
+    # Reset Train transforms
+    train_loader.dataset.transform = training_transforms
+    return result_dict

finetuning/map_function.py

@@ -0,0 +1,11 @@
+from finetuning.qsenn import finetune_qsenn
+from finetuning.sldd import finetune_sldd
+
+
+def finetune(key, model, train_loader, test_loader, log_dir, n_classes, seed, beta, optimization_schedule, per_class, n_features):
+    if key == 'sldd':
+        return finetune_sldd(model, train_loader, test_loader, log_dir, n_classes, seed, beta, optimization_schedule,per_class, n_features)
+    elif key == 'qsenn':
+        return finetune_qsenn(model, train_loader, test_loader, log_dir, n_classes, seed, beta, optimization_schedule,n_features,per_class, )
+    else:
+        raise ValueError(f"Unknown Finetuning key: {key}")

finetuning/qsenn.py

@@ -0,0 +1,30 @@
+import os
+
+import torch
+
+from finetuning.utils import train_n_epochs
+from sparsification.qsenn import compute_qsenn_feature_selection_and_assignment
+
+
+def finetune_qsenn(model, train_loader, test_loader, log_dir, n_classes, seed, beta, optimization_schedule ,n_features, n_per_class):
+    for iteration_epoch in range(4):
+        print(f"Starting iteration epoch {iteration_epoch}")
+        this_log_dir = log_dir / f"iteration_epoch_{iteration_epoch}"
+        this_log_dir.mkdir(parents=True, exist_ok=True)
+        feature_sel, sparse_layer,bias_sparse, current_mean, current_std = compute_qsenn_feature_selection_and_assignment(model, train_loader,
+                                                                                        test_loader,
+                                                                                        this_log_dir, n_classes, seed, n_features, n_per_class)
+        model.set_model_sldd(feature_sel, sparse_layer, current_mean, current_std, bias_sparse)
+        if os.path.exists(this_log_dir / "trained_model.pth"):
+            model.load_state_dict(torch.load(this_log_dir / "trained_model.pth"))
+            _ = optimization_schedule.get_params() # count up,  to have get correct lr
+            continue
+
+        model = train_n_epochs( model, beta, optimization_schedule, train_loader, test_loader)
+        torch.save(model.state_dict(), this_log_dir / "trained_model.pth")
+        print(f"Finished iteration epoch {iteration_epoch}")
+    return model
+
+
+
+

finetuning/sldd.py

@@ -0,0 +1,22 @@
+import numpy as np
+import torch
+
+from FeatureDiversityLoss import FeatureDiversityLoss
+from finetuning.utils import train_n_epochs
+from sparsification.glmBasedSparsification import compute_feature_selection_and_assignment
+from sparsification.sldd import compute_sldd_feature_selection_and_assignment
+from train import train, test
+from training.optim import get_optimizer
+
+
+
+
+def finetune_sldd(model, train_loader, test_loader, log_dir, n_classes, seed, beta, optimization_schedule,n_per_class, n_features, ):
+        feature_sel, weight, bias, mean, std = compute_sldd_feature_selection_and_assignment(model, train_loader,
+                                                                                        test_loader,
+                                                                                        log_dir, n_classes, seed,n_per_class, n_features)
+        model.set_model_sldd(feature_sel, weight, mean, std, bias)
+        model = train_n_epochs( model, beta, optimization_schedule, train_loader, test_loader)
+        return model
+
+