OneBitViT is a vision transformer model for image classification tasks. Args: image_size (int or tuple): The size of the input image. If an integer is provided, it is assumed to be a square image. patch_size (int or tuple): The size of each patch in the image. If an integer
| 80 | |
| 81 | |
| 82 | class OneBitViT(nn.Module): |
| 83 | """ |
| 84 | OneBitViT is a vision transformer model for image classification tasks. |
| 85 | |
| 86 | Args: |
| 87 | image_size (int or tuple): The size of the input image. If an integer is provided, it is assumed to be a square image. |
| 88 | patch_size (int or tuple): The size of each patch in the image. If an integer is provided, it is assumed to be a square patch. |
| 89 | num_classes (int): The number of output classes. |
| 90 | dim (int): The dimensionality of the token embeddings and the positional embeddings. |
| 91 | depth (int): The number of transformer layers. |
| 92 | heads (int): The number of attention heads in the transformer. |
| 93 | mlp_dim (int): The dimensionality of the feed-forward network in the transformer. |
| 94 | channels (int): The number of input channels in the image. Default is 3. |
| 95 | dim_head (int): The dimensionality of each attention head. Default is 64. |
| 96 | |
| 97 | Attributes: |
| 98 | to_patch_embedding (nn.Sequential): Sequential module for converting image patches to embeddings. |
| 99 | pos_embedding (torch.Tensor): Positional embeddings for the patches. |
| 100 | transformer (Transformer): Transformer module for processing the embeddings. |
| 101 | pool (str): Pooling method used to aggregate the patch embeddings. Default is "mean". |
| 102 | to_latent (nn.Identity): Identity module for converting the transformer output to the final latent representation. |
| 103 | linear_head (nn.LayerNorm): Layer normalization module for the final linear projection. |
| 104 | |
| 105 | Methods: |
| 106 | forward(img): Performs a forward pass through the OneBitViT model. |
| 107 | |
| 108 | """ |
| 109 | |
| 110 | def __init__( |
| 111 | self, |
| 112 | *, |
| 113 | image_size, |
| 114 | patch_size, |
| 115 | num_classes, |
| 116 | dim, |
| 117 | depth, |
| 118 | heads, |
| 119 | mlp_dim, |
| 120 | channels=3, |
| 121 | dim_head=64 |
| 122 | ): |
| 123 | super().__init__() |
| 124 | image_height, image_width = pair(image_size) |
| 125 | patch_height, patch_width = pair(patch_size) |
| 126 | |
| 127 | assert ( |
| 128 | image_height % patch_height == 0 and image_width % patch_width == 0 |
| 129 | ), "Image dimensions must be divisible by the patch size." |
| 130 | |
| 131 | patch_dim = channels * patch_height * patch_width |
| 132 | |
| 133 | self.to_patch_embedding = nn.Sequential( |
| 134 | Rearrange( |
| 135 | "b c (h p1) (w p2) -> b (h w) (p1 p2 c)", |
| 136 | p1=patch_height, |
| 137 | p2=patch_width, |
| 138 | ), |
| 139 | nn.LayerNorm(patch_dim), |