2495 Enhance PatchEmbeddingBlock

docs/source/networks.rst

@@ -80,7 +80,7 @@ Blocks
     :members:
 
 `SABlock Block`
-~~~~~~~~~~~~~~~~~
+~~~~~~~~~~~~~~~
 .. autoclass:: SABlock
     :members:
 
@@ -90,12 +90,12 @@ Blocks
     :members:
 
 `Transformer Block`
-~~~~~~~~~~~~~~~~~~~~~~~~
+~~~~~~~~~~~~~~~~~~~
 .. autoclass:: TransformerBlock
     :members:
 
 `UNETR Block`
-~~~~~~~~~~~~~~~~~~~~
+~~~~~~~~~~~~~
 .. autoclass:: UnetrBasicBlock
     :members:
 .. autoclass:: UnetrUpBlock
@@ -154,12 +154,12 @@ Blocks
     :members:
 
 `Registration Down Sample Block`
-~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
 .. autoclass:: RegistrationDownSampleBlock
     :members:
 
 `Registration Extraction Block`
-~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
 .. autoclass:: RegistrationExtractionBlock
     :members:
 
@@ -179,12 +179,12 @@ Blocks
     :members:
 
 `MLP Block`
-~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+~~~~~~~~~~~
 .. autoclass:: MLPBlock
     :members:
 
 `Patch Embedding Block`
-~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+~~~~~~~~~~~~~~~~~~~~~~~
 .. autoclass:: PatchEmbeddingBlock
     :members:
 
@@ -419,7 +419,7 @@ Nets
 .. autoclass:: unet
 
 `UNETR`
-~~~~~~~~~~~~~~~~~
+~~~~~~~
 .. autoclass:: UNETR
     :members:
 
@@ -436,7 +436,7 @@ Nets
   :members:
 
 `RegUNet`
-~~~~~~~~~~
+~~~~~~~~~
 .. autoclass:: RegUNet
   :members:
 
@@ -461,7 +461,7 @@ Nets
   :members:
 
 `ViT`
-~~~~~~
+~~~~~
 .. autoclass:: ViT
   :members:
 

monai/networks/blocks/patchembedding.py

@@ -18,7 +18,7 @@
 
 from monai.utils import optional_import
 
-einops, has_einops = optional_import("einops")
+Rearrange, _ = optional_import("einops.layers.torch", name="Rearrange")
 
 
 class PatchEmbeddingBlock(nn.Module):
@@ -30,12 +30,11 @@ class PatchEmbeddingBlock(nn.Module):
     def __init__(
         self,
         in_channels: int,
-        img_size: Union[int, Tuple[int, int, int]],
-        patch_size: Union[int, Tuple[int, int, int]],
+        img_size: Tuple[int, int, int],
+        patch_size: Tuple[int, int, int],
         hidden_size: int,
         num_heads: int,
-        pos_embed: Union[Tuple, str],  # type: ignore
-        classification: bool,
+        pos_embed: str,
         dropout_rate: float = 0.0,
     ) -> None:
         """
@@ -46,7 +45,6 @@ def __init__(
             hidden_size: dimension of hidden layer.
             num_heads: number of attention heads.
             pos_embed: position embedding layer type.
-            classification: bool argument to determine if classification is used.
             dropout_rate: faction of the input units to drop.
 
         """
@@ -59,39 +57,35 @@ def __init__(
         if hidden_size % num_heads != 0:
             raise AssertionError("hidden size should be divisible by num_heads.")
 
-        if img_size < patch_size:  # type: ignore
-            raise AssertionError("patch_size should be smaller than img_size.")
+        for m, p in zip(img_size, patch_size):
+            if m < p:
+                raise AssertionError("patch_size should be smaller than img_size.")
 
         if pos_embed not in ["conv", "perceptron"]:
             raise KeyError(f"Position embedding layer of type {pos_embed} is not supported.")
 
         if pos_embed == "perceptron":
-            if img_size[0] % patch_size[0] != 0:  # type: ignore
+            if img_size[0] % patch_size[0] != 0:
                 raise AssertionError("img_size should be divisible by patch_size for perceptron patch embedding.")
 
-        if has_einops:  # type: ignore
-            from einops.layers.torch import Rearrange  # type: ignore
-
-            self.Rearrange = Rearrange  # type: ignore
-        else:
-            raise ValueError('"Requires einops.')
-
         self.n_patches = (
-            (img_size[0] // patch_size[0]) * (img_size[1] // patch_size[1]) * (img_size[2] // patch_size[2])  # type: ignore
+            (img_size[0] // patch_size[0]) * (img_size[1] // patch_size[1]) * (img_size[2] // patch_size[2])
         )
-        self.patch_dim = in_channels * patch_size[0] * patch_size[1] * patch_size[2]  # type: ignore
+        self.patch_dim = in_channels * patch_size[0] * patch_size[1] * patch_size[2]
+
         self.pos_embed = pos_embed
+        self.patch_embeddings: Union[nn.Conv3d, nn.Sequential]
         if self.pos_embed == "conv":
             self.patch_embeddings = nn.Conv3d(
-                in_channels=in_channels, out_channels=hidden_size, kernel_size=patch_size, stride=patch_size  # type: ignore
+                in_channels=in_channels, out_channels=hidden_size, kernel_size=patch_size, stride=patch_size
             )
         elif self.pos_embed == "perceptron":
-            self.patch_embeddings = nn.Sequential(  # type: ignore
-                self.Rearrange(
+            self.patch_embeddings = nn.Sequential(
+                Rearrange(
                     "b c (h p1) (w p2) (d p3)-> b (h w d) (p1 p2 p3 c)",
-                    p1=patch_size[0],  # type: ignore
-                    p2=patch_size[1],  # type: ignore
-                    p3=patch_size[2],  # type: ignore
+                    p1=patch_size[0],
+                    p2=patch_size[1],
+                    p3=patch_size[2],
                 ),
                 nn.Linear(self.patch_dim, hidden_size),
             )

monai/networks/nets/unetr.py

@@ -28,12 +28,12 @@ def __init__(
         self,
         in_channels: int,
         out_channels: int,
-        img_size: Tuple,  # type: ignore
+        img_size: Tuple[int, int, int],
         feature_size: int,
         hidden_size: int,
         mlp_dim: int,
         num_heads: int,
-        pos_embed: Union[Tuple, str],
+        pos_embed: str,
         norm_name: Union[Tuple, str],
         conv_block: bool = False,
         res_block: bool = False,
@@ -70,9 +70,9 @@ def __init__(
         self.num_layers = 12
         self.patch_size = (16, 16, 16)
         self.feat_size = (
-            img_size[0] // self.patch_size[0],  # type: ignore
-            img_size[1] // self.patch_size[1],  # type: ignore
-            img_size[2] // self.patch_size[2],  # type: ignore
+            img_size[0] // self.patch_size[0],
+            img_size[1] // self.patch_size[1],
+            img_size[2] // self.patch_size[2],
         )
         self.hidden_size = hidden_size
         self.classification = False

monai/networks/nets/vit.py

@@ -10,7 +10,7 @@
 # limitations under the License.
 
 
-from typing import Tuple, Union
+from typing import Tuple
 
 import torch.nn as nn
 
@@ -27,13 +27,13 @@ class ViT(nn.Module):
     def __init__(
         self,
         in_channels: int,
-        img_size: Tuple,  # type: ignore
-        patch_size: Tuple,  # type: ignore
+        img_size: Tuple[int, int, int],
+        patch_size: Tuple[int, int, int],
         hidden_size: int,
         mlp_dim: int,
         num_layers: int,
         num_heads: int,
-        pos_embed: Union[Tuple, str],
+        pos_embed: str,
         classification: bool,
         num_classes: int = 2,
         dropout_rate: float = 0.0,
@@ -62,18 +62,15 @@ def __init__(
         if hidden_size % num_heads != 0:
             raise AssertionError("hidden size should be divisible by num_heads.")
 
-        if img_size < patch_size:
-            raise AssertionError("patch_size should be smaller than img_size.")
-
         if pos_embed not in ["conv", "perceptron"]:
             raise KeyError(f"Position embedding layer of type {pos_embed} is not supported.")
 
         self.classification = classification
         self.patch_embedding = PatchEmbeddingBlock(
-            in_channels, img_size, patch_size, hidden_size, num_heads, pos_embed, classification, dropout_rate  # type: ignore
+            in_channels, img_size, patch_size, hidden_size, num_heads, pos_embed, dropout_rate
         )
         self.blocks = nn.ModuleList(
-            [TransformerBlock(hidden_size, mlp_dim, num_heads, dropout_rate) for i in range(num_layers)]  # type: ignore
+            [TransformerBlock(hidden_size, mlp_dim, num_heads, dropout_rate) for i in range(num_layers)]
         )
         self.norm = nn.LayerNorm(hidden_size)
         if self.classification:

tests/test_patchembedding.py

@@ -43,7 +43,6 @@
                                         "hidden_size": hidden_size,
                                         "num_heads": num_heads,
                                         "pos_embed": pos_embed,
-                                        "classification": classification,
                                         "dropout_rate": dropout_rate,
                                     },
                                     (2, in_channels, img_size, *([img_size] * 2)),
@@ -70,7 +69,6 @@ def test_ill_arg(self):
                 hidden_size=128,
                 num_heads=12,
                 pos_embed="conv",
-                classification=False,
                 dropout_rate=5.0,
             )
 
@@ -82,7 +80,6 @@ def test_ill_arg(self):
                 hidden_size=512,
                 num_heads=8,
                 pos_embed="perceptron",
-                classification=False,
                 dropout_rate=0.3,
             )
 
@@ -94,7 +91,6 @@ def test_ill_arg(self):
                 hidden_size=512,
                 num_heads=14,
                 pos_embed="conv",
-                classification=False,
                 dropout_rate=0.3,
             )
 
@@ -106,7 +102,6 @@ def test_ill_arg(self):
                 hidden_size=768,
                 num_heads=8,
                 pos_embed="perceptron",
-                classification=False,
                 dropout_rate=0.3,
             )
 
@@ -118,7 +113,6 @@ def test_ill_arg(self):
                 hidden_size=768,
                 num_heads=12,
                 pos_embed="perc",
-                classification=False,
                 dropout_rate=0.3,
             )