| 207 | |
| 208 | |
| 209 | class VisionTransformer(nn.Module): |
| 210 | |
| 211 | def __init__(self, |
| 212 | image_size=224, |
| 213 | patch_size=16, |
| 214 | dim=768, |
| 215 | mlp_ratio=4, |
| 216 | out_dim=512, |
| 217 | num_heads=12, |
| 218 | num_layers=12, |
| 219 | pool_type='token', |
| 220 | pre_norm=True, |
| 221 | post_norm=False, |
| 222 | activation='quick_gelu', |
| 223 | attn_dropout=0.0, |
| 224 | proj_dropout=0.0, |
| 225 | embedding_dropout=0.0, |
| 226 | norm_eps=1e-5): |
| 227 | if image_size % patch_size != 0: |
| 228 | print( |
| 229 | '[WARNING] image_size is not divisible by patch_size', |
| 230 | flush=True) |
| 231 | assert pool_type in ('token', 'token_fc', 'attn_pool') |
| 232 | out_dim = out_dim or dim |
| 233 | super().__init__() |
| 234 | self.image_size = image_size |
| 235 | self.patch_size = patch_size |
| 236 | self.num_patches = (image_size // patch_size)**2 |
| 237 | self.dim = dim |
| 238 | self.mlp_ratio = mlp_ratio |
| 239 | self.out_dim = out_dim |
| 240 | self.num_heads = num_heads |
| 241 | self.num_layers = num_layers |
| 242 | self.pool_type = pool_type |
| 243 | self.post_norm = post_norm |
| 244 | self.norm_eps = norm_eps |
| 245 | |
| 246 | # embeddings |
| 247 | gain = 1.0 / math.sqrt(dim) |
| 248 | self.patch_embedding = nn.Conv2d( |
| 249 | 3, |
| 250 | dim, |
| 251 | kernel_size=patch_size, |
| 252 | stride=patch_size, |
| 253 | bias=not pre_norm) |
| 254 | if pool_type in ('token', 'token_fc'): |
| 255 | self.cls_embedding = nn.Parameter(gain * torch.randn(1, 1, dim)) |
| 256 | self.pos_embedding = nn.Parameter(gain * torch.randn( |
| 257 | 1, self.num_patches + |
| 258 | (1 if pool_type in ('token', 'token_fc') else 0), dim)) |
| 259 | self.dropout = nn.Dropout(embedding_dropout) |
| 260 | |
| 261 | # transformer |
| 262 | self.pre_norm = LayerNorm(dim, eps=norm_eps) if pre_norm else None |
| 263 | self.transformer = nn.Sequential(*[ |
| 264 | AttentionBlock(dim, mlp_ratio, num_heads, post_norm, False, |
| 265 | activation, attn_dropout, proj_dropout, norm_eps) |
| 266 | for _ in range(num_layers) |