| 384 | |
| 385 | |
| 386 | class VisionTransformer(nn.Module): |
| 387 | |
| 388 | def __init__(self, |
| 389 | image_size=224, |
| 390 | patch_size=16, |
| 391 | dim=768, |
| 392 | mlp_ratio=4, |
| 393 | out_dim=512, |
| 394 | num_heads=12, |
| 395 | num_layers=12, |
| 396 | pool_type='token', |
| 397 | pre_norm=True, |
| 398 | post_norm=False, |
| 399 | activation='quick_gelu', |
| 400 | attn_dropout=0.0, |
| 401 | proj_dropout=0.0, |
| 402 | embedding_dropout=0.0, |
| 403 | norm_eps=1e-5): |
| 404 | if image_size % patch_size != 0: |
| 405 | print( |
| 406 | '[WARNING] image_size is not divisible by patch_size', |
| 407 | flush=True) |
| 408 | assert pool_type in ('token', 'token_fc', 'attn_pool') |
| 409 | out_dim = out_dim or dim |
| 410 | super().__init__() |
| 411 | self.image_size = image_size |
| 412 | self.patch_size = patch_size |
| 413 | self.num_patches = (image_size // patch_size)**2 |
| 414 | self.dim = dim |
| 415 | self.mlp_ratio = mlp_ratio |
| 416 | self.out_dim = out_dim |
| 417 | self.num_heads = num_heads |
| 418 | self.num_layers = num_layers |
| 419 | self.pool_type = pool_type |
| 420 | self.post_norm = post_norm |
| 421 | self.norm_eps = norm_eps |
| 422 | |
| 423 | # embeddings |
| 424 | gain = 1.0 / math.sqrt(dim) |
| 425 | self.patch_embedding = nn.Conv2d( |
| 426 | 3, |
| 427 | dim, |
| 428 | kernel_size=patch_size, |
| 429 | stride=patch_size, |
| 430 | bias=not pre_norm) |
| 431 | if pool_type in ('token', 'token_fc'): |
| 432 | self.cls_embedding = nn.Parameter(gain * torch.randn(1, 1, dim)) |
| 433 | self.pos_embedding = nn.Parameter(gain * torch.randn( |
| 434 | 1, self.num_patches + |
| 435 | (1 if pool_type in ('token', 'token_fc') else 0), dim)) |
| 436 | self.dropout = nn.Dropout(embedding_dropout) |
| 437 | |
| 438 | # transformer |
| 439 | self.pre_norm = LayerNorm(dim, eps=norm_eps) if pre_norm else None |
| 440 | self.transformer = nn.Sequential(*[ |
| 441 | AttentionBlock(dim, mlp_ratio, num_heads, post_norm, False, |
| 442 | activation, attn_dropout, proj_dropout, norm_eps) |
| 443 | for _ in range(num_layers) |