Merge pull request #144 from cblessing24/mypy

Annotate Attention Readout

Author: christoph-blessing

.dockerignore

@@ -0,0 +1 @@
+.mypy_cache/

.github/workflows/mypy.yml

@@ -0,0 +1,12 @@
+name: Mypy
+
+on: [push, pull_request]
+
+jobs:
+  mypy:
+    runs-on: ubuntu-18.04
+    steps:
+    - uses: actions/checkout@v2
+    - uses: actions/setup-python@v2
+    - name: Check code with mypy
+      run: touch .env && docker-compose run mypy

.pre-commit-config.yaml

@@ -0,0 +1,10 @@
+repos:
+- repo: local
+  hooks:
+    - id: mypy
+      name: mypy
+      language: system
+      entry: docker-compose run --rm mypy
+      files: ^neuralpredictors/
+      types: [python]
+      pass_filenames: false

Dockerfile

@@ -1,10 +1,11 @@
-FROM sinzlab/pytorch:v3.8-torch1.7.0-cuda11.0-dj0.12.7
+FROM sinzlab/pytorch:v3.9-torch1.9.0-cuda11.1-dj0.12.7
 
 COPY . /src/neuralpredictors
+WORKDIR /src/neuralpredictors
 
 RUN python3 -m pip install --upgrade pip &&\
+    python3 -m pip install mypy==$(cat mypy_version.txt) &&\
     python3 -m pip install -e /src/neuralpredictors
 
-WORKDIR /src/neuralpredictors
 
 ENTRYPOINT ["python3"]

README.md

@@ -3,6 +3,7 @@
 ![Test](https://github.com/sinzlab/neuralpredictors/workflows/Test/badge.svg)
 [![codecov](https://codecov.io/gh/sinzlab/neuralpredictors/branch/main/graph/badge.svg)](https://codecov.io/gh/sinzlab/neuralpredictors)
 ![Black](https://github.com/sinzlab/neuralpredictors/workflows/Black/badge.svg)
+[![Mypy](https://github.com/sinzlab/neuralpredictors/actions/workflows/mypy.yml/badge.svg)](https://github.com/sinzlab/neuralpredictors/actions/workflows/mypy.yml)
 [![PyPI version](https://badge.fury.io/py/neuralpredictors.svg)](https://badge.fury.io/py/neuralpredictors)
 
 [Sinz Lab](https://sinzlab.org/) Neural System Identification Utilities for [PyTorch](https://pytorch.org/).

docker-compose.yml

@@ -18,3 +18,12 @@ services:
     build: .
     volumes:
       - .:/src/neuralpredictors
+  mypy:
+    build: .
+    volumes:
+      - .:/src/neuralpredictors
+      - mypy-cache:/src/neuralpredictors/.mypy_cache
+    entrypoint: ["/src/neuralpredictors/run_mypy.sh"]
+
+volumes:
+  mypy-cache:

mypy_files.txt

@@ -0,0 +1 @@
+neuralpredictors/layers/readouts/attention.py

mypy_version.txt

@@ -0,0 +1 @@
+0.910

neuralpredictors/layers/readouts/attention.py

@@ -1,88 +1,96 @@
+from typing import Any, Literal, Mapping, Optional, Tuple
+
 import torch
-from torch import nn
-from torch.nn import Parameter
 from torch.nn import functional as F
+from torch.nn import init
+from torch.nn.modules import ELU, BatchNorm2d, Conv2d, Module, Sequential
+from torch.nn.parameter import Parameter
+
 from .base import Readout
 
 
 class AttentionReadout(Readout):
     def __init__(
         self,
-        in_shape,
-        outdims,
-        bias,
-        init_noise=1e-3,
-        attention_kernel=1,
-        attention_layers=1,
-        mean_activity=None,
-        feature_reg_weight=1.0,
-        gamma_readout=None,  # depricated, use feature_reg_weight instead
-        **kwargs,
-    ):
+        in_shape: Tuple[int, int, int],
+        outdims: int,
+        bias: bool,
+        init_noise: float = 1e-3,
+        attention_kernel: int = 1,
+        attention_layers: int = 1,
+        mean_activity: Optional[Mapping[str, float]] = None,
+        feature_reg_weight: float = 1.0,
+        gamma_readout: Optional[float] = None,  # deprecated, use feature_reg_weight instead
+        **kwargs: Any,
+    ) -> None:
         super().__init__()
         self.in_shape = in_shape
         self.outdims = outdims
-        self.feature_reg_weight = self.resolve_deprecated_gamma_readout(feature_reg_weight, gamma_readout)
+        self.feature_reg_weight = self.resolve_deprecated_gamma_readout(feature_reg_weight, gamma_readout)  # type: ignore[no-untyped-call]
         self.mean_activity = mean_activity
         c, w, h = in_shape
         self.features = Parameter(torch.Tensor(self.outdims, c))
 
-        attention = nn.Sequential()
+        attention = Sequential()
         for i in range(attention_layers - 1):
             attention.add_module(
                 f"conv{i}",
-                nn.Conv2d(c, c, attention_kernel, padding=attention_kernel > 1),
+                Conv2d(c, c, attention_kernel, padding=attention_kernel > 1),
             )
-            attention.add_module(f"norm{i}", nn.BatchNorm2d(c))
-            attention.add_module(f"nonlin{i}", nn.ELU())
+            attention.add_module(f"norm{i}", BatchNorm2d(c))  # type: ignore[no-untyped-call]
+            attention.add_module(f"nonlin{i}", ELU())
         else:
             attention.add_module(
                 f"conv{attention_layers}",
-                nn.Conv2d(c, outdims, attention_kernel, padding=attention_kernel > 1),
+                Conv2d(c, outdims, attention_kernel, padding=attention_kernel > 1),
             )
         self.attention = attention
 
         self.init_noise = init_noise
         if bias:
-            bias = Parameter(torch.Tensor(self.outdims))
-            self.register_parameter("bias", bias)
+            bias_param = Parameter(torch.Tensor(self.outdims))
+            self.register_parameter("bias", bias_param)
         else:
             self.register_parameter("bias", None)
         self.initialize(mean_activity)
 
     @staticmethod
-    def init_conv(m):
-        if isinstance(m, nn.Conv2d):
-            nn.init.xavier_normal_(m.weight.data)
+    def init_conv(m: Module) -> None:
+        if isinstance(m, Conv2d):
+            init.xavier_normal_(m.weight.data)
             if m.bias is not None:
                 m.bias.data.fill_(0)
 
-    def initialize_attention(self):
+    def initialize_attention(self) -> None:
         self.apply(self.init_conv)
 
-    def initialize(self, mean_activity=None):
+    def initialize(self, mean_activity: Optional[Mapping[str, float]] = None) -> None:  # type: ignore[override]
         if mean_activity is None:
             mean_activity = self.mean_activity
         self.features.data.normal_(0, self.init_noise)
         if self.bias is not None:
-            self.initialize_bias(mean_activity=mean_activity)
+            self.initialize_bias(mean_activity=mean_activity)  # type: ignore[no-untyped-call]
         self.initialize_attention()
 
-    def feature_l1(self, reduction="sum", average=None):
-        return self.apply_reduction(self.features.abs(), reduction=reduction, average=average)
+    def feature_l1(
+        self, reduction: Literal["sum", "mean", None] = "sum", average: Optional[bool] = None
+    ) -> torch.Tensor:
+        return self.apply_reduction(self.features.abs(), reduction=reduction, average=average)  # type: ignore[no-untyped-call,no-any-return]
 
-    def regularizer(self, reduction="sum", average=None):
-        return self.feature_l1(reduction=reduction, average=average) * self.feature_reg_weight
+    def regularizer(
+        self, reduction: Literal["sum", "mean", None] = "sum", average: Optional[bool] = None
+    ) -> torch.Tensor:
+        return self.feature_l1(reduction=reduction, average=average) * self.feature_reg_weight  # type: ignore[no-any-return]
 
-    def forward(self, x, shift=None):
+    def forward(self, x: torch.Tensor, shift: Optional[Any] = None) -> torch.Tensor:
         attention = self.attention(x)
         b, c, w, h = attention.shape
         attention = F.softmax(attention.view(b, c, -1), dim=-1).view(b, c, w, h)
-        y = torch.einsum("bnwh,bcwh->bcn", attention, x)
-        y = torch.einsum("bcn,nc->bn", y, self.features)
+        y: torch.Tensor = torch.einsum("bnwh,bcwh->bcn", attention, x)  # type: ignore[attr-defined]
+        y = torch.einsum("bcn,nc->bn", y, self.features)  # type: ignore[attr-defined]
         if self.bias is not None:
             y = y + self.bias
         return y
 
-    def __repr__(self):
+    def __repr__(self) -> str:
         return self.__class__.__name__ + " (" + "{} x {} x {}".format(*self.in_shape) + " -> " + str(self.outdims) + ")"

pyproject.toml

@@ -4,3 +4,24 @@ line-length = 120
 [tool.coverage.run]
 branch = true
 source = ["neuralpredictors"]
+
+[tool.mypy]
+python_version = "3.8"
+files = "neuralpredictors"
+exclude = "old_\\w+\\.py$"
+strict = true
+disallow_untyped_calls = true
+disallow_untyped_defs = true
+disallow_incomplete_defs = true
+disallow_untyped_decorators = true
+
+[[tool.mypy.overrides]]
+module = [
+    "h5py",
+    "scipy.signal",
+    "scipy.special",
+    "skimage.transform",
+    "torchvision",
+    "tqdm"
+]
+ignore_missing_imports = true