commit from qixuan

.gitattributes

@@ -33,3 +33,4 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
 *.zip filter=lfs diff=lfs merge=lfs -text
 *.zst filter=lfs diff=lfs merge=lfs -text
 *tfevents* filter=lfs diff=lfs merge=lfs -text
+*.jpg filter=lfs diff=lfs merge=lfs -text

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)
+

ReadME.md

@@ -0,0 +1,138 @@
+# Q-SENN - Quantized Self-Explaining Neural Networks
+
+This repository contains the code for the AAAI 2024 paper 
+[*Q-SENN: Quantized Self-Explaining Neural Network*](https://ojs.aaai.org/index.php/AAAI/article/view/30145) by Thomas
+Norrenbrock , 
+Marco Rudolph,
+and Bodo Rosenhahn.
+Additonally, the SLDD-model from [*Take 5: 
+Interpretable Image Classification with a Handful of Features*](https://arxiv.org/pdf/2303.13166) (NeurIPS 
+Workshop) from the same authors is included.
+
+
+<p align="center">
+    <img width="400" height="400" src="fig/birds.png"> 
+</p>
+
+---
+Abstract:
+>Explanations in Computer Vision are often desired, but most Deep Neural Networks can only provide saliency maps with questionable faithfulness. Self-Explaining Neural Networks (SENN) extract interpretable concepts with fidelity, diversity, and grounding to combine them linearly for decision-making. While they can explain what was recognized, initial realizations lack accuracy and general applicability. We propose the Quantized-Self-Explaining Neural Network Q-SENN. Q-SENN satisfies or exceeds the desiderata of SENN while being applicable to more complex datasets and maintaining most or all of the accuracy of an uninterpretable baseline model, out-performing previous work in all considered metrics. Q-SENN describes the relationship between every class and feature as either positive, negative or neutral instead of an arbitrary number of possible relations, enforcing more binary human-friendly features. Since every class is assigned just 5 interpretable features on average, Q-SENN shows convincing local and global interpretability. Additionally, we propose a feature alignment method, capable of aligning learned features with human language-based concepts without additional supervision. Thus, what is learned can be more easily verbalized.
+
+
+
+
+---
+
+## Installation
+You will need the usual libaries for deep learning, e.g. pytorch, 
+torchvision, numpy, etc. Additionally, we use 
+[GLM-Saga](https://github.com/MadryLab/glm_saga) that can be installed via pip.
+In case you are lazy (or like to spend your time otherwise), a suitable 
+environment can be created using [Anaconda](https://www.anaconda.com/) and the 
+provided environment.yml file:
+```shell
+conda env create -f environment.yml 
+```
+
+## Data
+Supported datasets are:
+- [Cub2011](https://www.vision.caltech.edu/datasets/cub_200_2011/)
+- [StanfordCars](https://ai.stanford.edu/~jkrause/cars/car_dataset.html)
+- [TravelingBirds](https://worksheets.codalab.org/bundles/0x518829de2aa440c79cd9d75ef6669f27)
+- [ImageNet](https://www.image-net.org/)
+
+To use the data for training, the datasets have to be downloaded and put into the 
+respective folder under ~/tmp/datasets such that the final structure looks like
+
+```shell
+~/tmp/datasets
+├── CUB200
+│   └── CUB_200_2011
+│       ├── ...
+├── StanfordCars
+│   ├── stanford_cars
+│       ├── ...
+├── TravelingBirds
+│   ├── CUB_fixed
+│       ├── ...
+├── imagenet
+│   ├── ...
+```
+
+The default paths could be changed in the dataset_classes or for Imagenet in 
+get_data.py
+
+Note: 
+If cropped images, like for PIP-Net, ProtoPool, etc. are desired, 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, 
+default path is: PPCUB200  instead of CUB200 for Protopool. Using these images 
+can be set using an additional flag `--cropGT` introduced later.
+
+
+
+## Usage
+The code to create a Q-SENN model can be started from the file main.py.
+Available parameters are:
+- `--dataset`: The dataset to use. Default: Cub2011
+- `--arch`: The backbone to use. Default: resnet50
+- `--model_type`: The model type to use. Default: qsenn
+- `--seed`: The seed to use. Default: None
+- `--do_dense`: Whether to train the dense model. Default: True
+- `--cropGT`: Whether to crop CUB/TravelingBirds based on GT Boundaries. Default: False
+- `--n_features`: How many features to select. Default: 50
+- `--n_per_class`: How many features to assign to each class. Default: 5
+- `--img_size`: Image size. Default: 448
+- `--reduced_strides`: Whether to use reduced strides for resnets. Default: False
+
+
+For Example the next command will start the creation of Q-SENN with resnet50 on 
+StanfordCars using the default arguments in the paper.
+```shell
+python main.py --dataset StanfordCars
+```
+
+**Note:**
+All experiments on ImageNet in the paper skipped the dense training from 
+scratch on ImageNet. The pretrained models are used directly.
+This can be replicated with the argument --do-dense False.
+## Citations
+Please cite this work as:\
+Q-SENN
+```bibtex
+@inproceedings{norrenbrock2024q,
+  title={Q-senn: Quantized self-explaining neural networks},
+  author={Norrenbrock, Thomas and Rudolph, Marco and Rosenhahn, Bodo},
+  booktitle={Proceedings of the AAAI Conference on Artificial Intelligence},
+  volume={38},
+  number={19},
+  pages={21482--21491},
+  year={2024}
+}
+```
+SLDD-Model
+```bibtex
+@inproceedings{norrenbrocktake,
+  title={Take 5: Interpretable Image Classification with a Handful of Features},
+  author={Norrenbrock, Thomas and Rudolph, Marco and Rosenhahn, Bodo},
+  year={2022},
+  booktitle={Progress and Challenges in Building Trustworthy Embodied AI}
+}
+```
+## Pretrained Model
+One pretrained model for Q-SENN on CUB can be obtained via this link: https://drive.google.com/drive/folders/1agWqKhcWOVWueV4Fzaowr80lQroCJFYn?usp=drive_link
+## Acknowledgement
+This work was supported by the Federal Ministry of Education and Research (BMBF), Germany under the AI service center KISSKI (grant no. 01IS22093C) and the Deutsche Forschungsgemeinschaft (DFG) under Germany’s Excellence Strategy within the Cluster of Excellence PhoenixD (EXC 2122).
+This work was partially supported by Intel Corporation and by the German Federal Ministry
+of the Environment, Nature Conservation, Nuclear Safety
+and Consumer Protection (GreenAutoML4FAS project no.
+67KI32007A). 
+
+The work was done at the Leibniz University Hannover and published at AAAI 2024.
+
+<p align="center">
+    <img width="100" height="100" src="fig/AutoML4FAS_Logo.jpeg"> 
+    <img width="300" height="100" src="fig/Bund.png">
+    <img width="300" height="100" src="fig/LUH.png">
+</p>

app.py

@@ -0,0 +1,87 @@
+import gradio as gr
+from visualization_gary import*
+
+# 定义模式选择操作，加载不同的界面
+def choose_mode(selected_mode):
+    if selected_mode == "get interpretable Result ":
+        return gr.update(visible=True), gr.update(visible=False), gr.update(visible=False), "in Mode 1"
+    elif selected_mode == "To give Feedbacks":
+        return gr.update(visible=False), gr.update(visible=True), gr.update(visible=False), "in Mode 2"
+    else:
+        return gr.update(visible=False), gr.update(visible=False), gr.update(visible=False), "please choose a mode"
+
+# clear if reupload
+def reset_state_mod1():
+    return None, None, "State has been reset."
+
+def reset_state_mod2():
+    return None, None,None,gr.update(interactive=True),gr.update(interactive=True),gr.update(interactive=True),gr.update(interactive=True),gr.update(interactive=True)
+
+# 主界面，包含模式选择和每种模式的界面
+with gr.Blocks() as demo:
+    gr.Markdown("# please choose a mode")
+
+    # 模式选择器
+    mode_selector = gr.Radio(["get interpretable Result ", "To give Feedbacks"], label="Mode Selection")
+
+    # 模式1界面
+    with gr.Column(visible=False) as mode1_ui:
+        gr.Markdown("### Please keep the object in the center of your image, click the Button to get interpretebale Result")
+        mode1_button = gr.Button("get interpretable Classification")
+        mode1_input_img=gr.Image()
+        mode1_output_img= gr.Image()
+        mode1_output_txt=gr.Markdown()
+
+        #clear if reupload
+        mode1_state = gr.State()
+        mode1_input_img.upload(fn=reset_state_mod1, outputs=[mode1_output_img, mode1_state, mode1_output_txt])
+        #clear if reupload
+
+
+        mode1_button.click(fn=direct_inference,inputs=mode1_input_img, outputs=[mode1_output_img,mode1_output_txt])
+
+    # 模式2界面
+    with gr.Column(visible=False) as mode2_ui:
+        gr.Markdown("### Please keep the object in the center of your image, click the Button 'Get some interpretable Features' to get options")
+        image_input=gr.Image()
+        gallery_output = gr.Gallery(label="Initial Label")
+        text_output=gr.Markdown()
+        but_generate=gr.Button("Get some interpretable Features")
+        but_feedback_A=gr.Button("A")
+        but_feedback_B=gr.Button("B")
+        but_feedback_C=gr.Button("C")
+        but_feedback_D=gr.Button("D")
+
+        key_op = gr.State()
+        A= gr.State("A")
+        B= gr.State("B")
+        C= gr.State("C")
+        D= gr.State("D")
+        
+        #clear if reupload
+        image_input.upload(
+            fn=reset_state_mod2, 
+            outputs=[gallery_output, key_op, text_output,but_generate,but_feedback_A,but_feedback_B,but_feedback_C,but_feedback_D]
+        )
+
+        #clear if reupload
+
+
+
+        but_generate.click(fn=get_features_on_interface, inputs=image_input, outputs=[gallery_output,key_op,text_output,but_generate])
+
+        but_feedback_A.click(fn=post_next_image, inputs=[A,key_op], outputs=[text_output,but_feedback_A,but_feedback_B,but_feedback_C,but_feedback_D])
+        but_feedback_B.click(fn=post_next_image, inputs=[B,key_op], outputs=[text_output,but_feedback_A,but_feedback_B,but_feedback_C,but_feedback_D])
+        but_feedback_C.click(fn=post_next_image, inputs=[C,key_op], outputs=[text_output,but_feedback_A,but_feedback_B,but_feedback_C,but_feedback_D])
+        but_feedback_D.click(fn=post_next_image, inputs=[D,key_op], outputs=[text_output,but_feedback_A,but_feedback_B,but_feedback_C,but_feedback_D])
+        # but_feedback_B.click(fn=post_next_image, inputs=image_list, outputs=[image_list,image_output,text_output])
+
+
+
+    # 状态输出
+    status_output = gr.Textbox(label="Status")
+
+    # 选择器点击事件绑定
+    mode_selector.change(choose_mode, inputs=mode_selector, outputs=[mode1_ui, mode2_ui, status_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/qixuan/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