Source code for monai.networks.nets.vitautoenc

# Copyright (c) MONAI Consortium
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# See the License for the specific language governing permissions and
# limitations under the License.

from __future__ import annotations

from import Sequence

import torch
import torch.nn as nn

from monai.networks.blocks.patchembedding import PatchEmbeddingBlock
from monai.networks.blocks.transformerblock import TransformerBlock
from monai.networks.layers import Conv
from monai.utils import ensure_tuple_rep

__all__ = ["ViTAutoEnc"]

[docs]class ViTAutoEnc(nn.Module): """ Vision Transformer (ViT), based on: "Dosovitskiy et al., An Image is Worth 16x16 Words: Transformers for Image Recognition at Scale <>" Modified to also give same dimension outputs as the input size of the image """
[docs] def __init__( self, in_channels: int, img_size: Sequence[int] | int, patch_size: Sequence[int] | int, out_channels: int = 1, deconv_chns: int = 16, hidden_size: int = 768, mlp_dim: int = 3072, num_layers: int = 12, num_heads: int = 12, pos_embed: str = "conv", dropout_rate: float = 0.0, spatial_dims: int = 3, qkv_bias: bool = False, save_attn: bool = False, ) -> None: """ Args: in_channels: dimension of input channels or the number of channels for input. img_size: dimension of input image. patch_size: dimension of patch size out_channels: number of output channels. Defaults to 1. deconv_chns: number of channels for the deconvolution layers. Defaults to 16. hidden_size: dimension of hidden layer. Defaults to 768. mlp_dim: dimension of feedforward layer. Defaults to 3072. num_layers: number of transformer blocks. Defaults to 12. num_heads: number of attention heads. Defaults to 12. pos_embed: position embedding layer type. Defaults to "conv". dropout_rate: faction of the input units to drop. Defaults to 0.0. spatial_dims: number of spatial dimensions. Defaults to 3. qkv_bias: apply bias to the qkv linear layer in self attention block. Defaults to False. save_attn: to make accessible the attention in self attention block. Defaults to False. Defaults to False. Examples:: # for single channel input with image size of (96,96,96), conv position embedding and segmentation backbone # It will provide an output of same size as that of the input >>> net = ViTAutoEnc(in_channels=1, patch_size=(16,16,16), img_size=(96,96,96), pos_embed='conv') # for 3-channel with image size of (128,128,128), output will be same size as of input >>> net = ViTAutoEnc(in_channels=3, patch_size=(16,16,16), img_size=(128,128,128), pos_embed='conv') """ super().__init__() self.patch_size = ensure_tuple_rep(patch_size, spatial_dims) self.spatial_dims = spatial_dims self.patch_embedding = PatchEmbeddingBlock( in_channels=in_channels, img_size=img_size, patch_size=patch_size, hidden_size=hidden_size, num_heads=num_heads, pos_embed=pos_embed, dropout_rate=dropout_rate, spatial_dims=self.spatial_dims, ) self.blocks = nn.ModuleList( [ TransformerBlock(hidden_size, mlp_dim, num_heads, dropout_rate, qkv_bias, save_attn) for i in range(num_layers) ] ) self.norm = nn.LayerNorm(hidden_size) new_patch_size = [4] * self.spatial_dims conv_trans = Conv[Conv.CONVTRANS, self.spatial_dims] # self.conv3d_transpose* is to be compatible with existing 3d model weights. self.conv3d_transpose = conv_trans(hidden_size, deconv_chns, kernel_size=new_patch_size, stride=new_patch_size) self.conv3d_transpose_1 = conv_trans( in_channels=deconv_chns, out_channels=out_channels, kernel_size=new_patch_size, stride=new_patch_size )
[docs] def forward(self, x): """ Args: x: input tensor must have isotropic spatial dimensions, such as ``[batch_size, channels, sp_size, sp_size[, sp_size]]``. """ spatial_size = x.shape[2:] x = self.patch_embedding(x) hidden_states_out = [] for blk in self.blocks: x = blk(x) hidden_states_out.append(x) x = self.norm(x) x = x.transpose(1, 2) d = [s // p for s, p in zip(spatial_size, self.patch_size)] x = torch.reshape(x, [x.shape[0], x.shape[1], *d]) x = self.conv3d_transpose(x) x = self.conv3d_transpose_1(x) return x, hidden_states_out