Vision Transformer with support for patch or hybrid CNN input stage
| 326 | |
| 327 | |
| 328 | class ViT(nn.Module): |
| 329 | """ Vision Transformer with support for patch or hybrid CNN input stage |
| 330 | """ |
| 331 | |
| 332 | def __init__(self, img_size=224, patch_size=16, in_chans=3, num_classes=80, embed_dim=768, depth=12, |
| 333 | num_heads=12, mlp_ratio=4., qkv_bias=False, |
| 334 | drop_path_rate=0., norm_layer=None, window=True, |
| 335 | use_abs_pos_emb=False, interval=3, bn_group=None, test_pos_mode=False, |
| 336 | task_sp_list=(), neck_sp_list=(), learnable_pos=False, rel_pos_spatial=False, lms_checkpoint_train=False, |
| 337 | prompt=None, pad_attn_mask=False, freeze_iters=0, |
| 338 | act_layer='GELU', pre_ln=False, mask_input=False, ending_norm=True, |
| 339 | round_padding=False, compat=False, use_cls_token=False): |
| 340 | super().__init__() |
| 341 | self.pad_attn_mask = pad_attn_mask # only effective for detection task input w/ NestedTensor wrapping |
| 342 | self.lms_checkpoint_train = lms_checkpoint_train |
| 343 | self.use_cls_token = use_cls_token |
| 344 | self.task_sp_list = task_sp_list |
| 345 | self.neck_sp_list = neck_sp_list |
| 346 | self.freeze_iters = freeze_iters |
| 347 | self.mask_input = mask_input |
| 348 | self.ending_norm = ending_norm |
| 349 | self.round_padding = round_padding |
| 350 | |
| 351 | global COMPAT |
| 352 | COMPAT = compat |
| 353 | |
| 354 | norm_layer = norm_layer or partial(nn.LayerNorm, eps=1e-6) |
| 355 | self.num_classes = num_classes |
| 356 | self.num_features = self.embed_dim = embed_dim # num_features for consistency with other models |
| 357 | |
| 358 | self.patch_embed = PatchEmbed( |
| 359 | img_size=img_size, patch_size=patch_size, in_chans=in_chans, embed_dim=embed_dim) |
| 360 | |
| 361 | num_patches = self.patch_embed.num_patches |
| 362 | if use_abs_pos_emb: |
| 363 | if self.use_cls_token: |
| 364 | self.cls_token = nn.Parameter(torch.zeros(1, 1, embed_dim)) |
| 365 | self.pos_embed = nn.Parameter(torch.zeros(1, num_patches + 1, embed_dim), requires_grad=learnable_pos) |
| 366 | trunc_normal_(self.cls_token, std=.02) |
| 367 | trunc_normal_(self.pos_embed, std=.02) |
| 368 | else: |
| 369 | self.pos_embed = nn.Parameter(torch.zeros(1, num_patches, embed_dim), requires_grad=learnable_pos) |
| 370 | pos_embed = get_2d_sincos_pos_embed(self.pos_embed.shape[-1], self.patch_embed.patch_shape, cls_token=False) |
| 371 | self.pos_embed.data.copy_(torch.from_numpy(pos_embed).float().unsqueeze(0)) |
| 372 | else: |
| 373 | raise |
| 374 | |
| 375 | dpr = [x.item() for x in torch.linspace(0, drop_path_rate, depth)] # stochastic depth decay rule |
| 376 | |
| 377 | self.blocks = nn.ModuleList() |
| 378 | for i in range(depth): |
| 379 | block = Block( |
| 380 | dim=embed_dim, num_heads=num_heads, mlp_ratio=mlp_ratio, qkv_bias=qkv_bias, |
| 381 | drop_path=dpr[i], norm_layer=norm_layer, |
| 382 | window_size=(14, 14) if ((i + 1) % interval != 0) else self.patch_embed.patch_shape, |
| 383 | window=((i + 1) % interval != 0) if window else False, |
| 384 | rel_pos_spatial=rel_pos_spatial, prompt=prompt, |
| 385 | act_layer=QuickGELU if act_layer == 'QuickGELU' else nn.GELU |
no outgoing calls
no test coverage detected