| |
|
|
| |
| |
| |
| |
| |
| |
| |
| |
| |
| |
|
|
|
|
| import math |
| from typing import Sequence, Tuple, Union |
|
|
| import torch |
| import torch.nn as nn |
| import torch.nn.functional as F |
| from monai.utils import optional_import |
|
|
| Rearrange, _ = optional_import("einops.layers.torch", name="Rearrange") |
|
|
|
|
| class PatchEmbeddingBlock(nn.Module): |
| """ |
| A patch embedding block, based on: "Dosovitskiy et al., |
| An Image is Worth 16x16 Words: Transformers for Image Recognition at Scale <https://arxiv.org/abs/2010.11929>" |
| """ |
|
|
| def __init__( |
| self, |
| in_channels: int, |
| img_size: Tuple[int, int, int], |
| patch_size: Tuple[int, int, int], |
| hidden_size: int, |
| num_heads: int, |
| pos_embed: str, |
| dropout_rate: float = 0.0, |
| ) -> None: |
| """ |
| Args: |
| in_channels: dimension of input channels. |
| img_size: dimension of input image. |
| patch_size: dimension of patch size. |
| hidden_size: dimension of hidden layer. |
| num_heads: number of attention heads. |
| pos_embed: position embedding layer type. |
| dropout_rate: faction of the input units to drop. |
| |
| """ |
|
|
| super().__init__() |
|
|
| if not (0 <= dropout_rate <= 1): |
| raise AssertionError("dropout_rate should be between 0 and 1.") |
|
|
| if hidden_size % num_heads != 0: |
| raise AssertionError("hidden size should be divisible by num_heads.") |
|
|
| 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: |
| raise AssertionError("img_size should be divisible by patch_size for perceptron patch embedding.") |
|
|
| self.n_patches = ( |
| (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] |
|
|
| 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 |
| ) |
| elif self.pos_embed == "perceptron": |
| 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], |
| p2=patch_size[1], |
| p3=patch_size[2], |
| ), |
| nn.Linear(self.patch_dim, hidden_size), |
| ) |
| self.position_embeddings = nn.Parameter(torch.zeros(1, self.n_patches, hidden_size)) |
| self.cls_token = nn.Parameter(torch.zeros(1, 1, hidden_size)) |
| self.dropout = nn.Dropout(dropout_rate) |
| self.trunc_normal_(self.position_embeddings, mean=0.0, std=0.02, a=-2.0, b=2.0) |
| self.apply(self._init_weights) |
|
|
| def _init_weights(self, m): |
| if isinstance(m, nn.Linear): |
| self.trunc_normal_(m.weight, mean=0.0, std=0.02, a=-2.0, b=2.0) |
| if isinstance(m, nn.Linear) and m.bias is not None: |
| nn.init.constant_(m.bias, 0) |
| elif isinstance(m, nn.LayerNorm): |
| nn.init.constant_(m.bias, 0) |
| nn.init.constant_(m.weight, 1.0) |
|
|
| def trunc_normal_(self, tensor, mean, std, a, b): |
| |
| |
| def norm_cdf(x): |
| return (1.0 + math.erf(x / math.sqrt(2.0))) / 2.0 |
|
|
| with torch.no_grad(): |
| l = norm_cdf((a - mean) / std) |
| u = norm_cdf((b - mean) / std) |
| tensor.uniform_(2 * l - 1, 2 * u - 1) |
| tensor.erfinv_() |
| tensor.mul_(std * math.sqrt(2.0)) |
| tensor.add_(mean) |
| tensor.clamp_(min=a, max=b) |
| return tensor |
|
|
| def forward(self, x): |
| if self.pos_embed == "conv": |
| x = self.patch_embeddings(x) |
| x = x.flatten(2) |
| x = x.transpose(-1, -2) |
| elif self.pos_embed == "perceptron": |
| x = self.patch_embeddings(x) |
| embeddings = x + self.position_embeddings |
| embeddings = self.dropout(embeddings) |
| return embeddings |
|
|
|
|
| class PatchEmbed3D(nn.Module): |
| """Video to Patch Embedding. |
| |
| Args: |
| patch_size (int): Patch token size. Default: (2,4,4). |
| in_chans (int): Number of input video channels. Default: 3. |
| embed_dim (int): Number of linear projection output channels. Default: 96. |
| norm_layer (nn.Module, optional): Normalization layer. Default: None |
| """ |
|
|
| def __init__( |
| self, |
| img_size: Sequence[int] = (96, 96, 96), |
| patch_size=(4, 4, 4), |
| in_chans: int = 1, |
| embed_dim: int = 96, |
| norm_layer=None, |
| ): |
| super().__init__() |
| self.patch_size = patch_size |
|
|
| self.in_chans = in_chans |
| self.embed_dim = embed_dim |
|
|
| self.grid_size = (img_size[0] // patch_size[0], img_size[1] // patch_size[1], img_size[2] // patch_size[2]) |
| self.num_patches = self.grid_size[0] * self.grid_size[1] * self.grid_size[2] |
|
|
| self.proj = nn.Conv3d(in_chans, embed_dim, kernel_size=patch_size, stride=patch_size) |
|
|
| if norm_layer is not None: |
| self.norm = norm_layer(embed_dim) |
| else: |
| self.norm = None |
|
|
| def forward(self, x): |
| """Forward function.""" |
| |
| _, _, d, h, w = x.size() |
| if w % self.patch_size[2] != 0: |
| x = F.pad(x, (0, self.patch_size[2] - w % self.patch_size[2])) |
| if h % self.patch_size[1] != 0: |
| x = F.pad(x, (0, 0, 0, self.patch_size[1] - h % self.patch_size[1])) |
| if d % self.patch_size[0] != 0: |
| x = F.pad(x, (0, 0, 0, 0, 0, self.patch_size[0] - d % self.patch_size[0])) |
|
|
| x = self.proj(x) |
| if self.norm is not None: |
| d, wh, ww = x.size(2), x.size(3), x.size(4) |
| x = x.flatten(2).transpose(1, 2) |
| x = self.norm(x) |
| x = x.transpose(1, 2).view(-1, self.embed_dim, d, wh, ww) |
| |
|
|
| return x |
|
|
