Add mean pooling layer (#187)

CHANGELOG.md

@@ -17,6 +17,7 @@ the section headers (Added/Changed/...) and incrementing the package version.
 the environment file since it is necessary for the augmentations.
 - ([#178](https://github.com/microsoft/hi-ml/pull/178)) Add runner script for running ML experiments.
 - ([#181](https://github.com/microsoft/hi-ml/pull/181)) Add computational pathology tools in hi-ml-histopathology folder.
+- ([#187](https://github.com/microsoft/hi-ml/pull/187)) Add mean pooling layer for MIL.
 
 ### Changed
 

hi-ml/src/health_ml/networks/layers/attention_layers.py

@@ -6,11 +6,27 @@
 Created using the original DeepMIL paper and code from Ilse et al., 2018
 https://github.com/AMLab-Amsterdam/AttentionDeepMIL (MIT License)
 """
-from typing import Tuple
+from typing import Any, Tuple
 from torch import nn, Tensor, transpose, mm
+import torch
 import torch.nn.functional as F
 
 
+class MeanPoolingLayer(nn.Module):
+    """Mean pooling returns uniform weights and the average feature vector over the first axis"""
+
+    # args/kwargs added here for compatibility with parametrised pooling modules
+    def __init__(self, *args: Any, **kwargs: Any) -> None:
+        super().__init__()
+
+    def forward(self, features: Tensor) -> Tuple[Tensor, Tensor]:
+        num_instances = features.shape[0]
+        A = torch.full((1, num_instances), 1. / num_instances)
+        M = features.mean(dim=0)
+        M = M.view(1, -1)
+        return (A, M)
+
+
 class AttentionLayer(nn.Module):
     """ AttentionLayer: Simple attention layer
     Requires size of input L, hidden D, and attention layers K (default K=1)

hi-ml/testhiml/testhiml/test_attentionlayers.py

@@ -1,31 +1,43 @@
 import pytest
 from typing import Type, Union
 
-from torch import rand, sum, allclose, ones_like
+from torch import nn, rand, sum, allclose, ones_like
 
-from health_ml.networks.layers.attention_layers import AttentionLayer, GatedAttentionLayer
+from health_ml.networks.layers.attention_layers import (AttentionLayer, GatedAttentionLayer,
+                                                        MeanPoolingLayer)
+
+
+def _test_attention_layer(attentionlayer: nn.Module, dim_in: int, dim_att: int,
+                          batch_size: int,) -> None:
+    features = rand(batch_size, dim_in)                   # N x L x 1 x 1
+    attn_weights, output_features = attentionlayer(features)
+    assert attn_weights.shape == (dim_att, batch_size)          # K x N
+    assert output_features.shape == (dim_att, dim_in)           # K x L
+    assert ((attn_weights >= 0) & (attn_weights <= 1)).all()
+
+    row_sums = sum(attn_weights, dim=1, keepdim=True)
+    assert allclose(row_sums, ones_like(row_sums))
+
+    pooled_features = attn_weights @ features.flatten(start_dim=1)
+    assert allclose(pooled_features, output_features)
 
 
 @pytest.mark.parametrize("dim_in", [1, 3])
 @pytest.mark.parametrize("dim_hid", [1, 4])
 @pytest.mark.parametrize("dim_att", [1, 5])
 @pytest.mark.parametrize("batch_size", [1, 7])
 @pytest.mark.parametrize('attention_layer_cls', [AttentionLayer, GatedAttentionLayer])
-def test_attentionlayer(dim_in: int,
-                        dim_hid: int,
-                        dim_att: int,
-                        batch_size: int,
+def test_attentionlayer(dim_in: int, dim_hid: int, dim_att: int, batch_size: int,
                         attention_layer_cls: Type[Union[AttentionLayer, GatedAttentionLayer]]) -> None:
-
     attentionlayer = attention_layer_cls(
         input_dims=dim_in,
         hidden_dims=dim_hid,
         attention_dims=dim_att
     )
+    _test_attention_layer(attentionlayer, dim_in, dim_att, batch_size)
 
-    features = rand(batch_size, dim_in, 1, 1)                   # N x L x 1 x 1
-    attn_weights, output_features = attentionlayer(features)
-    assert attn_weights.shape == (dim_att, batch_size)          # K x N
-    assert output_features.shape == (dim_att, dim_in)           # K x L
-    row_sums = sum(attn_weights, dim=1, keepdim=True)
-    assert allclose(row_sums, ones_like(row_sums))
+
+@pytest.mark.parametrize("dim_in", [1, 3])
+@pytest.mark.parametrize("batch_size", [1, 7])
+def test_mean_pooling(dim_in: int, batch_size: int,) -> None:
+    _test_attention_layer(MeanPoolingLayer(), dim_in=dim_in, dim_att=1, batch_size=batch_size)