| 15 | |
| 16 | # This class and its supporting functions below lightly adapted from the ViTDet backbone available at: https://github.com/facebookresearch/detectron2/blob/main/detectron2/modeling/backbone/vit.py # noqa |
| 17 | class ImageEncoderViT(nn.Module): |
| 18 | def __init__( |
| 19 | self, |
| 20 | img_size: int = 1024, |
| 21 | patch_size: int = 16, |
| 22 | in_chans: int = 3, |
| 23 | embed_dim: int = 768, |
| 24 | depth: int = 12, |
| 25 | num_heads: int = 12, |
| 26 | mlp_ratio: float = 4.0, |
| 27 | out_chans: int = 256, |
| 28 | qkv_bias: bool = True, |
| 29 | norm_layer: Type[nn.Module] = nn.LayerNorm, |
| 30 | act_layer: Type[nn.Module] = nn.GELU, |
| 31 | use_abs_pos: bool = True, |
| 32 | use_rel_pos: bool = False, |
| 33 | rel_pos_zero_init: bool = True, |
| 34 | window_size: int = 0, |
| 35 | global_attn_indexes: Tuple[int, ...] = (), |
| 36 | ) -> None: |
| 37 | """ |
| 38 | Args: |
| 39 | img_size (int): Input image size. |
| 40 | patch_size (int): Patch size. |
| 41 | in_chans (int): Number of input image channels. |
| 42 | embed_dim (int): Patch embedding dimension. |
| 43 | depth (int): Depth of ViT. |
| 44 | num_heads (int): Number of attention heads in each ViT block. |
| 45 | mlp_ratio (float): Ratio of mlp hidden dim to embedding dim. |
| 46 | qkv_bias (bool): If True, add a learnable bias to query, key, value. |
| 47 | norm_layer (nn.Module): Normalization layer. |
| 48 | act_layer (nn.Module): Activation layer. |
| 49 | use_abs_pos (bool): If True, use absolute positional embeddings. |
| 50 | use_rel_pos (bool): If True, add relative positional embeddings to the attention map. |
| 51 | rel_pos_zero_init (bool): If True, zero initialize relative positional parameters. |
| 52 | window_size (int): Window size for window attention blocks. |
| 53 | global_attn_indexes (list): Indexes for blocks using global attention. |
| 54 | """ |
| 55 | super().__init__() |
| 56 | self.img_size = img_size |
| 57 | |
| 58 | self.patch_embed = PatchEmbed( |
| 59 | kernel_size=(patch_size, patch_size), |
| 60 | stride=(patch_size, patch_size), |
| 61 | in_chans=in_chans, |
| 62 | embed_dim=embed_dim, |
| 63 | ) |
| 64 | |
| 65 | self.pos_embed: Optional[nn.Parameter] = None |
| 66 | if use_abs_pos: |
| 67 | # Initialize absolute positional embedding with pretrain image size. |
| 68 | self.pos_embed = nn.Parameter( |
| 69 | torch.zeros(1, img_size // patch_size, img_size // patch_size, embed_dim) |
| 70 | ) |
| 71 | |
| 72 | self.blocks = nn.ModuleList() |
| 73 | for i in range(depth): |
| 74 | block = Block( |