Reenable and fix mypy-torch errors

- Argument "{}" to "{}" has incompatible type "{}"; expected "{}"
- Incompatible types in assignment (expression has type "{}", variable has type "{}")
- Module "{}" is not valid as a type
- Module has no attribute "{}"
- Name '{}' is not defined
- No overload variant of "{}" of "{}" matches argument types

Author: Benjamin Gorman

monai/data/csv_saver.py

@@ -77,6 +77,7 @@ def save(self, data: Union[torch.Tensor, np.ndarray], meta_data: Optional[dict]
         self._data_index += 1
         if torch.is_tensor(data):
             data = data.detach().cpu().numpy()
+        assert isinstance(data, np.ndarray)
         self._cache_dict[save_key] = data.astype(np.float32)
 
     def save_batch(self, batch_data: Union[torch.Tensor, np.ndarray], meta_data: Optional[dict] = None) -> None:

monai/data/dataloader.py

@@ -63,7 +63,7 @@ def __init__(
         timeout: float = 0.0,
         multiprocessing_context: Optional[Callable] = None,
     ) -> None:
-        super().__init__(
+        super().__init__(  # type: ignore # No overload variant matches argument types
             dataset=dataset,
             batch_size=batch_size,
             shuffle=shuffle,

monai/data/utils.py

@@ -452,7 +452,7 @@ def compute_importance_map(
     patch_size,
     mode: Union[BlendMode, str] = BlendMode.CONSTANT,
     sigma_scale: float = 0.125,
-    device: Optional[Union[torch.device, str]] = None,
+    device: Optional[torch.device] = None,
 ):
     """Get importance map for different weight modes.
 

monai/engines/multi_gpu_supervised_trainer.py

@@ -15,6 +15,7 @@
     import ignite.metrics
 
 import torch
+from torch.optim.optimizer import Optimizer
 
 from monai.utils import exact_version, optional_import
 from monai.engines.utils import get_devices_spec
@@ -34,7 +35,7 @@ def _default_eval_transform(x, y, y_pred):
 
 def create_multigpu_supervised_trainer(
     net: torch.nn.Module,
-    optimizer: torch.optim.Optimizer,
+    optimizer: Optimizer,
     loss_fn,
     devices=None,
     non_blocking: bool = False,

monai/engines/trainer.py

@@ -16,6 +16,7 @@
     import ignite.metrics
 
 import torch
+from torch.optim.optimizer import Optimizer
 from torch.utils.data import DataLoader
 
 from monai.inferers import Inferer, SimpleInferer
@@ -78,7 +79,7 @@ def __init__(
         max_epochs: int,
         train_data_loader: DataLoader,
         network,
-        optimizer: torch.optim.Optimizer,
+        optimizer: Optimizer,
         loss_function,
         prepare_batch: Callable = default_prepare_batch,
         iteration_update: Optional[Callable] = None,

monai/handlers/lr_schedule_handler.py

@@ -9,15 +9,15 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from typing import Callable, Optional, TYPE_CHECKING
+from typing import Callable, Optional, Union, TYPE_CHECKING
 
 if TYPE_CHECKING:
     import ignite.engine
 
-import torch
-
 import logging
 
+from torch.optim.lr_scheduler import _LRScheduler, ReduceLROnPlateau
+
 from monai.utils import ensure_tuple, exact_version, optional_import
 
 Events, _ = optional_import("ignite.engine", "0.3.0", exact_version, "Events")
@@ -30,7 +30,7 @@ class LrScheduleHandler:
 
     def __init__(
         self,
-        lr_scheduler: torch.optim.lr_scheduler,
+        lr_scheduler: Union[_LRScheduler, ReduceLROnPlateau],
         print_lr: bool = True,
         name: Optional[str] = None,
         epoch_level: bool = True,
@@ -73,4 +73,6 @@ def __call__(self, engine: "ignite.engine.Engine") -> None:
         args = ensure_tuple(self.step_transform(engine))
         self.lr_scheduler.step(*args)
         if self.print_lr:
-            self.logger.info(f"Current learning rate: {self.lr_scheduler._last_lr[0]}")
+            self.logger.info(
+                f"Current learning rate: {self.lr_scheduler._last_lr[0]}"  # type: ignore # Module has no attribute
+            )

monai/handlers/mean_dice.py

@@ -9,7 +9,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from typing import Callable, Optional, Sequence, Union
+from typing import Callable, Optional, Sequence
 
 import torch
 
@@ -72,13 +72,13 @@ def reset(self) -> None:
         self._num_examples = 0
 
     @reinit__is_reduced
-    def update(self, output: Sequence[Union[torch.Tensor, dict]]) -> None:
+    def update(self, output: Sequence[torch.Tensor]) -> None:
         if not len(output) == 2:
             raise ValueError("MeanDice metric can only support y_pred and y.")
         y_pred, y = output
         score = self.dice(y_pred, y)
         assert self.dice.not_nans is not None
-        not_nans = self.dice.not_nans.item()
+        not_nans = int(self.dice.not_nans.item())
 
         # add all items in current batch
         self._sum += score.item() * not_nans

monai/handlers/roc_auc.py

@@ -63,8 +63,8 @@ def __init__(
         self.average: Average = Average(average)
 
     def reset(self) -> None:
-        self._predictions: List[float] = []
-        self._targets: List[float] = []
+        self._predictions: List[torch.Tensor] = []
+        self._targets: List[torch.Tensor] = []
 
     def update(self, output: Sequence[torch.Tensor]) -> None:
         y_pred, y = output

monai/losses/dice.py

@@ -66,7 +66,7 @@ def __init__(
             ValueError: sigmoid=True and softmax=True are not compatible.
 
         """
-        super().__init__(reduction=LossReduction(reduction))
+        super().__init__(reduction=LossReduction(reduction).value)
 
         if sigmoid and softmax:
             raise ValueError("sigmoid=True and softmax=True are not compatible.")
@@ -133,11 +133,11 @@ def forward(self, input: torch.Tensor, target: torch.Tensor, smooth: float = 1e-
 
         f = 1.0 - (2.0 * intersection + smooth) / (denominator + smooth)
 
-        if self.reduction == LossReduction.MEAN:
+        if self.reduction == LossReduction.MEAN.value:
             f = torch.mean(f)  # the batch and channel average
-        elif self.reduction == LossReduction.SUM:
+        elif self.reduction == LossReduction.SUM.value:
             f = torch.sum(f)  # sum over the batch and channel dims
-        elif self.reduction == LossReduction.NONE:
+        elif self.reduction == LossReduction.NONE.value:
             pass  # returns [N, n_classes] losses
         else:
             raise ValueError(f"reduction={self.reduction} is invalid.")
@@ -219,7 +219,7 @@ def __init__(
             ValueError: sigmoid=True and softmax=True are not compatible.
 
         """
-        super().__init__(reduction=LossReduction(reduction))
+        super().__init__(reduction=LossReduction(reduction).value)
 
         self.include_background = include_background
         self.to_onehot_y = to_onehot_y
@@ -286,11 +286,11 @@ def forward(self, input: torch.Tensor, target: torch.Tensor, smooth: float = 1e-
 
         f = 1.0 - (2.0 * (intersection * w).sum(1) + smooth) / ((denominator * w).sum(1) + smooth)
 
-        if self.reduction == LossReduction.MEAN:
+        if self.reduction == LossReduction.MEAN.value:
             f = torch.mean(f)  # the batch and channel average
-        elif self.reduction == LossReduction.SUM:
+        elif self.reduction == LossReduction.SUM.value:
             f = torch.sum(f)  # sum over the batch and channel dims
-        elif self.reduction == LossReduction.NONE:
+        elif self.reduction == LossReduction.NONE.value:
             pass  # returns [N, n_classes] losses
         else:
             raise ValueError(f"reduction={self.reduction} is invalid.")

monai/losses/focal_loss.py

@@ -57,7 +57,7 @@ def __init__(
                 fl(pred, grnd)
 
         """
-        super(FocalLoss, self).__init__(weight=weight, reduction=LossReduction(reduction))
+        super(FocalLoss, self).__init__(weight=weight, reduction=LossReduction(reduction).value)
         self.gamma = gamma
         self.weight: Optional[torch.Tensor]
 
@@ -129,10 +129,10 @@ def forward(self, input: torch.Tensor, target: torch.Tensor):
         else:
             loss = torch.mean(-weight * t * logpt, dim=-1)  # N,C
 
-        if self.reduction == LossReduction.SUM:
+        if self.reduction == LossReduction.SUM.value:
             return loss.sum()
-        if self.reduction == LossReduction.NONE:
+        if self.reduction == LossReduction.NONE.value:
             return loss
-        if self.reduction == LossReduction.MEAN:
+        if self.reduction == LossReduction.MEAN.value:
             return loss.mean()
         raise ValueError(f"reduction={self.reduction} is invalid.")

monai/losses/tversky.py

@@ -63,7 +63,7 @@ def __init__(
 
         """
 
-        super().__init__(reduction=LossReduction(reduction))
+        super().__init__(reduction=LossReduction(reduction).value)
         self.include_background = include_background
         self.to_onehot_y = to_onehot_y
 
@@ -125,10 +125,10 @@ def forward(self, input: torch.Tensor, target: torch.Tensor, smooth: float = 1e-
 
         score = 1.0 - numerator / denominator
 
-        if self.reduction == LossReduction.SUM:
+        if self.reduction == LossReduction.SUM.value:
             return score.sum()  # sum over the batch and channel dims
-        if self.reduction == LossReduction.NONE:
+        if self.reduction == LossReduction.NONE.value:
             return score  # returns [N, n_classes] losses
-        if self.reduction == LossReduction.MEAN:
+        if self.reduction == LossReduction.MEAN.value:
             return score.mean()
         raise ValueError(f"reduction={self.reduction} is invalid.")

monai/metrics/rocauc.py

@@ -9,7 +9,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from typing import Union
+from typing import cast, Union
 
 import warnings
 
@@ -34,7 +34,7 @@ def _calculate(y: torch.Tensor, y_pred: torch.Tensor):
     nneg = auc = tmp_pos = tmp_neg = 0.0
 
     for i in range(n):
-        y_i = y[i]
+        y_i = cast(float, y[i])
         if i + 1 < n and y_pred[i] == y_pred[i + 1]:
             tmp_pos += y_i
             tmp_neg += 1 - y_i

monai/networks/blocks/upsample.py

@@ -48,7 +48,7 @@ def __init__(
             align_corners: set the align_corners parameter of `torch.nn.Upsample`. Defaults to True.
         """
         super().__init__()
-        scale_factor = ensure_tuple_rep(scale_factor, spatial_dims)
+        scale_factor_ = ensure_tuple_rep(scale_factor, spatial_dims)
         if not out_channels:
             out_channels = in_channels
         if not with_conv:
@@ -58,11 +58,11 @@ def __init__(
                 mode = linear_mode[spatial_dims - 1]
             self.upsample = nn.Sequential(
                 Conv[Conv.CONV, spatial_dims](in_channels=in_channels, out_channels=out_channels, kernel_size=1),
-                nn.Upsample(scale_factor=scale_factor, mode=mode.value, align_corners=align_corners),
+                nn.Upsample(scale_factor=scale_factor_, mode=mode.value, align_corners=align_corners),
             )
         else:
             self.upsample = Conv[Conv.CONVTRANS, spatial_dims](
-                in_channels=in_channels, out_channels=out_channels, kernel_size=scale_factor, stride=scale_factor
+                in_channels=in_channels, out_channels=out_channels, kernel_size=scale_factor_, stride=scale_factor_
             )
 
     def forward(self, x: torch.Tensor) -> torch.Tensor:

monai/networks/layers/factories.py

@@ -204,8 +204,8 @@ def batch_factory(dim: int):
 Act.add_factory_callable("gelu", lambda: nn.modules.GELU)
 Act.add_factory_callable("sigmoid", lambda: nn.modules.Sigmoid)
 Act.add_factory_callable("tanh", lambda: nn.modules.Tanh)
-Act.add_factory_callable("softmax", lambda: nn.modules.SoftMax)
-Act.add_factory_callable("logsoftmax", lambda: nn.modules.LogSoftMax)
+Act.add_factory_callable("softmax", lambda: nn.modules.Softmax)
+Act.add_factory_callable("logsoftmax", lambda: nn.modules.LogSoftmax)
 
 
 @Conv.factory_function("conv")

monai/networks/layers/spatial_transforms.py

@@ -147,7 +147,7 @@ def forward(self, src, theta, spatial_size: Optional[Union[Sequence[int], int]]
                 )
             )
 
-        grid = nn.functional.affine_grid(theta=theta[:, :sr], size=dst_size, align_corners=self.align_corners)
+        grid = nn.functional.affine_grid(theta=theta[:, :sr], size=list(dst_size), align_corners=self.align_corners)
         dst = nn.functional.grid_sample(
             input=src.contiguous(),
             grid=grid,

monai/networks/nets/densenet.py

@@ -144,7 +144,7 @@ def __init__(
                 [
                     ("relu", nn.ReLU(inplace=True)),
                     ("norm", avg_pool_type(1)),
-                    ("flatten", nn.Flatten(1)),
+                    ("flatten", nn.Flatten(1)),  # type: ignore # Module has no attribute
                     ("class", nn.Linear(in_channels, out_channels)),
                 ]
             )

monai/networks/nets/generator.py

@@ -81,7 +81,7 @@ def __init__(
         self.dropout = dropout
         self.bias = bias
 
-        self.flatten = nn.Flatten()
+        self.flatten = nn.Flatten()  # type: ignore # Module has no attribute
         self.linear = nn.Linear(int(np.prod(self.latent_shape)), int(np.prod(start_shape)))
         self.reshape = Reshape(*start_shape)
         self.conv = nn.Sequential()
@@ -102,6 +102,8 @@ def _get_layer(self, in_channels: int, out_channels: int, strides, is_last: bool
         is True this is the final layer and is not expected to include activation and normalization layers.
         """
 
+        layer: Union[Convolution, nn.Sequential]
+
         layer = Convolution(
             in_channels=in_channels,
             strides=strides,

monai/networks/nets/regressor.py

@@ -96,6 +96,8 @@ def _get_layer(self, in_channels: int, out_channels: int, strides, is_last: bool
         is True this is the final layer and is not expected to include activation and normalization layers.
         """
 
+        layer: Union[ResidualUnit, Convolution]
+
         if self.num_res_units > 0:
             layer = ResidualUnit(
                 subunits=self.num_res_units,

monai/transforms/post/array.py

@@ -344,5 +344,5 @@ def __call__(self, img: torch.Tensor):
         else:
             raise RuntimeError("the dimensions of img should be 4 or 5.")
 
-        torch.clamp_(contour_img, min=0.0, max=1.0)
+        contour_img.clamp_(min=0.0, max=1.0)
         return contour_img

monai/transforms/spatial/array.py

@@ -13,7 +13,7 @@
 https://github.com/Project-MONAI/MONAI/wiki/MONAI_Design
 """
 
-from typing import List, Optional, Sequence, Tuple, Union
+from typing import Callable, List, Optional, Sequence, Tuple, Union
 
 import warnings
 
@@ -47,6 +47,8 @@
 
 nib, _ = optional_import("nibabel")
 
+_torch_interp: Callable
+
 if get_torch_version_tuple() >= (1, 5):
     # additional argument since torch 1.5 (to avoid warnings)
     def _torch_interp(**kwargs):

monai/utils/misc.py

@@ -197,9 +197,9 @@ def set_determinism(
     """
     if seed is None:
         # cast to 32 bit seed for CUDA
-        seed_ = torch.default_generator.seed() % (np.iinfo(np.int32).max + 1)
-        if not torch.cuda._is_in_bad_fork():
-            torch.cuda.manual_seed_all(seed_)
+        seed_ = torch.default_generator.seed() % (np.iinfo(np.int32).max + 1)  # type: ignore # Module has no attribute
+        if not torch.cuda._is_in_bad_fork():  # type: ignore # Module has no attribute
+            torch.cuda.manual_seed_all(seed_)  # type: ignore # Module has no attribute
     else:
         torch.manual_seed(seed)
 
@@ -214,7 +214,7 @@ def set_determinism(
             func(seed)
 
     if seed is not None:
-        torch.backends.cudnn.deterministic = True
-        torch.backends.cudnn.benchmark = False
+        torch.backends.cudnn.deterministic = True  # type: ignore # Module has no attribute
+        torch.backends.cudnn.benchmark = False  # type: ignore # Module has no attribute
     else:
-        torch.backends.cudnn.deterministic = False
+        torch.backends.cudnn.deterministic = False  # type: ignore # Module has no attribute

setup.cfg

@@ -91,20 +91,7 @@ warn_no_return=True
 # lint-style cleanliness for typing needs to be disabled; returns more errors
 # than the full run.
 warn_redundant_casts=True
-warn_unused_ignores=False
-
-
-[mypy-torch.*]
-# --- Temporary disabling to allow smaller PRs
-# Subsequent fixes necessary to make this tool pass are not yet
-# merged into the master branch, but requested as separate
-# commits.  Can not enable until all outstanding typehint
-# PR's are approved
-## Temporarily disable all mypy warnings
-ignore_errors = True
-# ^^^^^^ Temporary disabling to allow smaller PRs
-follow_imports = skip
-follow_imports_for_stubs = True
+warn_unused_ignores=True
 
 [mypy-versioneer]
 # Always ignore any type issues in the versioneer.py file