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