Vision Transformer with support for patch or hybrid CNN input stage
| 137 | |
| 138 | |
| 139 | class VisionTransformer(nn.Module): |
| 140 | """ Vision Transformer with support for patch or hybrid CNN input stage |
| 141 | """ |
| 142 | def __init__(self, img_size=224, patch_size=16, in_chans=3, num_classes=1000, embed_dim=768, depth=12, |
| 143 | num_heads=12, mlp_ratio=4., qkv_bias=False, qk_scale=None, drop_rate=0., attn_drop_rate=0., |
| 144 | drop_path_rate=0., hybrid_backbone=None, norm_layer=nn.LayerNorm, is_distill=False): |
| 145 | super().__init__() |
| 146 | |
| 147 | if isinstance(img_size,tuple): |
| 148 | self.img_size = img_size |
| 149 | else: |
| 150 | self.img_size = to_2tuple(img_size) |
| 151 | |
| 152 | self.depth = depth |
| 153 | self.patch_size = patch_size |
| 154 | self.in_chans = in_chans |
| 155 | self.embed_dim = embed_dim |
| 156 | self.num_classes = num_classes |
| 157 | self.num_features = self.embed_dim = embed_dim # num_features for consistency with other models |
| 158 | |
| 159 | if hybrid_backbone is not None: |
| 160 | self.patch_embed = HybridEmbed( |
| 161 | hybrid_backbone, img_size=img_size, in_chans=in_chans, embed_dim=embed_dim) |
| 162 | else: |
| 163 | self.patch_embed = PatchEmbed( |
| 164 | img_size=img_size, patch_size=patch_size, in_chans=in_chans, embed_dim=embed_dim) |
| 165 | self.num_patches = self.patch_embed.num_patches |
| 166 | |
| 167 | self.cls_token = nn.Parameter(torch.zeros(1, 1, embed_dim)) |
| 168 | if is_distill: |
| 169 | self.pos_embed = nn.Parameter(torch.zeros(1, self.num_patches + 2, embed_dim)) |
| 170 | else: |
| 171 | self.pos_embed = nn.Parameter(torch.zeros(1, self.num_patches + 1, embed_dim)) |
| 172 | self.pos_drop = nn.Dropout(p=drop_rate) |
| 173 | |
| 174 | dpr = [x.item() for x in torch.linspace(0, drop_path_rate, depth)] # stochastic depth decay rule |
| 175 | self.blocks = nn.ModuleList([ |
| 176 | Block( |
| 177 | dim=embed_dim, num_heads=num_heads, mlp_ratio=mlp_ratio, qkv_bias=qkv_bias, qk_scale=qk_scale, |
| 178 | drop=drop_rate, attn_drop=attn_drop_rate, drop_path=dpr[i], norm_layer=norm_layer) |
| 179 | for i in range(depth)]) |
| 180 | self.norm = norm_layer(embed_dim) |
| 181 | self.head = nn.Linear(embed_dim, num_classes) |
| 182 | self.det_token_num = 0 |
| 183 | self.use_checkpoint = False |
| 184 | |
| 185 | # NOTE as per official impl, we could have a pre-logits representation dense layer + tanh here |
| 186 | #self.repr = nn.Linear(embed_dim, representation_size) |
| 187 | #self.repr_act = nn.Tanh() |
| 188 | |
| 189 | # Classifier head |
| 190 | # self.head = nn.Linear(embed_dim, num_classes) if num_classes > 0 else nn.Identity() |
| 191 | |
| 192 | trunc_normal_(self.pos_embed, std=.02) |
| 193 | trunc_normal_(self.cls_token, std=.02) |
| 194 | self.apply(self._init_weights) |
| 195 | |
| 196 |