Vision Transformer with support for patch or hybrid CNN input stage
| 363 | |
| 364 | |
| 365 | class ViT(nn.Module): |
| 366 | """ Vision Transformer with support for patch or hybrid CNN input stage |
| 367 | """ |
| 368 | |
| 369 | def __init__(self, img_size=224, patch_size=16, in_chans=3, num_classes=80, embed_dim=768, depth=12, |
| 370 | num_heads=12, mlp_ratio=4., qkv_bias=False, drop_rate=0., attn_drop_rate=0., |
| 371 | drop_path_rate=0., hybrid_backbone=None, norm_layer=None, init_values=None, window=False, |
| 372 | use_abs_pos_emb=False, interval=3, pretrained=None, bn_group=None, proj_padding=True, test_pos_mode=False): |
| 373 | super().__init__() |
| 374 | self.proj_padding = proj_padding |
| 375 | norm_layer = norm_layer or partial(nn.LayerNorm, eps=1e-6) |
| 376 | self.num_classes = num_classes |
| 377 | self.num_features = self.embed_dim = embed_dim # num_features for consistency with other models |
| 378 | |
| 379 | if hybrid_backbone is not None: |
| 380 | self.patch_embed = HybridEmbed( |
| 381 | hybrid_backbone, img_size=img_size, in_chans=in_chans, embed_dim=embed_dim) |
| 382 | else: |
| 383 | self.patch_embed = PatchEmbed( |
| 384 | img_size=img_size, patch_size=patch_size, in_chans=in_chans, embed_dim=embed_dim, proj_padding=self.proj_padding) |
| 385 | |
| 386 | num_patches = self.patch_embed.num_patches |
| 387 | |
| 388 | if use_abs_pos_emb: |
| 389 | self.pos_embed = nn.Parameter(torch.zeros(1, num_patches + 1, embed_dim)) |
| 390 | else: |
| 391 | raise |
| 392 | |
| 393 | self.pos_drop = nn.Dropout(p=drop_rate) |
| 394 | |
| 395 | dpr = [x.item() for x in torch.linspace(0, drop_path_rate, depth)] # stochastic depth decay rule |
| 396 | |
| 397 | if window: |
| 398 | self.blocks = nn.ModuleList([ |
| 399 | Block( |
| 400 | dim=embed_dim, num_heads=num_heads, mlp_ratio=mlp_ratio, qkv_bias=qkv_bias, |
| 401 | drop=drop_rate, attn_drop=attn_drop_rate, drop_path=dpr[i], norm_layer=norm_layer, |
| 402 | init_values=init_values, |
| 403 | window_size=(14, 14) if ((i + 1) % interval != 0) else self.patch_embed.patch_shape, |
| 404 | window=((i + 1) % interval != 0)) |
| 405 | for i in range(depth)]) |
| 406 | else: |
| 407 | self.blocks = nn.ModuleList([ |
| 408 | Block( |
| 409 | dim=embed_dim, num_heads=num_heads, mlp_ratio=mlp_ratio, qkv_bias=qkv_bias, |
| 410 | drop=drop_rate, attn_drop=attn_drop_rate, drop_path=dpr[i], norm_layer=norm_layer, |
| 411 | init_values=init_values, |
| 412 | window_size=(14, 14) if ((i + 1) % interval != 0) else self.patch_embed.patch_shape, |
| 413 | window=False) |
| 414 | for i in range(depth)]) |
| 415 | |
| 416 | self.norm = norm_layer(embed_dim) |
| 417 | |
| 418 | self.apply(self._init_weights) |
| 419 | self.fix_init_weight() |
| 420 | self.pretrained = pretrained |
| 421 | self.test_pos_mode = test_pos_mode |
| 422 |
no outgoing calls
no test coverage detected