Add novel heads: Optimal Transport (det), Info Bottleneck (seg), Harmonic Depth (depth)

heads/__init__.py

@@ -8,6 +8,7 @@ from .wavelet.head import Wavelet
 from .patch_attention.head import PatchAttention
 from .graph_crf.head import GraphCRF
 from .hypercolumn_linear.head import HypercolumnLinear
+from .info_bottleneck.head import InfoBottleneck
 
 REGISTRY = {
     "linear_probe": LinearProbe,
@@ -18,6 +19,7 @@ REGISTRY = {
     "patch_attention": PatchAttention,
     "graph_crf": GraphCRF,
     "hypercolumn_linear": HypercolumnLinear,
+    "info_bottleneck": InfoBottleneck,
 }
 
 ALL_NAMES = list(REGISTRY.keys())

heads/info_bottleneck/__init__.py

@@ -0,0 +1 @@
+from .head import InfoBottleneck

heads/info_bottleneck/head.py

@@ -0,0 +1,24 @@
+"""Information Bottleneck: project to d dimensions, classify from the compressed representation.
+
+The bottleneck dimension d is the minimum sufficient statistic for segmentation.
+If d=8 works for 150-class ADE20K, the frozen features have at most 8 independent
+directions relevant to semantic segmentation. That is a statement about the
+backbone's feature geometry.
+"""
+
+import torch.nn as nn
+
+
+class InfoBottleneck(nn.Module):
+    name = "info_bottleneck"
+    needs_intermediates = False
+
+    def __init__(self, feat_dim=768, num_classes=150, bottleneck_dim=8):
+        super().__init__()
+        self.bottleneck_dim = bottleneck_dim
+        self.compress = nn.Conv2d(feat_dim, bottleneck_dim, 1, bias=False)
+        self.classify = nn.Conv2d(bottleneck_dim, num_classes, 1)
+
+    def forward(self, spatial, inter=None):
+        z = self.compress(spatial)
+        return self.classify(z)