diff --git a/docs/source/networks.rst b/docs/source/networks.rst index f70fc4f024..a5ce86287a 100644 --- a/docs/source/networks.rst +++ b/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: diff --git a/monai/networks/blocks/patchembedding.py b/monai/networks/blocks/patchembedding.py index 6b80fdcc40..1f312e9126 100644 --- a/monai/networks/blocks/patchembedding.py +++ b/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), ) diff --git a/monai/networks/nets/unetr.py b/monai/networks/nets/unetr.py index 4e8b68b43e..5bd77f97d8 100644 --- a/monai/networks/nets/unetr.py +++ b/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 diff --git a/monai/networks/nets/vit.py b/monai/networks/nets/vit.py index 1a43414aca..a90e8d069e 100644 --- a/monai/networks/nets/vit.py +++ b/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: diff --git a/tests/test_patchembedding.py b/tests/test_patchembedding.py index c19c970c2b..5283153880 100644 --- a/tests/test_patchembedding.py +++ b/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, )