(self, img_size, patch_size, stride, base_dims, depth, heads,
mlp_ratio, num_classes=1000, in_chans=3, distilled=False,
attn_drop_rate=.0, drop_rate=.0, drop_path_rate=.0)
| 150 | - https://arxiv.org/abs/2103.16302 |
| 151 | """ |
| 152 | def __init__(self, img_size, patch_size, stride, base_dims, depth, heads, |
| 153 | mlp_ratio, num_classes=1000, in_chans=3, distilled=False, |
| 154 | attn_drop_rate=.0, drop_rate=.0, drop_path_rate=.0): |
| 155 | super(PoolingVisionTransformer, self).__init__() |
| 156 | |
| 157 | padding = 0 |
| 158 | img_size = to_2tuple(img_size) |
| 159 | patch_size = to_2tuple(patch_size) |
| 160 | height = math.floor((img_size[0] + 2 * padding - patch_size[0]) / stride + 1) |
| 161 | width = math.floor((img_size[1] + 2 * padding - patch_size[1]) / stride + 1) |
| 162 | |
| 163 | self.base_dims = base_dims |
| 164 | self.heads = heads |
| 165 | self.num_classes = num_classes |
| 166 | self.num_tokens = 2 if distilled else 1 |
| 167 | |
| 168 | self.patch_size = patch_size |
| 169 | self.pos_embed = nn.Parameter(torch.randn(1, base_dims[0] * heads[0], height, width)) |
| 170 | self.patch_embed = ConvEmbedding(in_chans, base_dims[0] * heads[0], patch_size, stride, padding) |
| 171 | |
| 172 | self.cls_token = nn.Parameter(torch.randn(1, self.num_tokens, base_dims[0] * heads[0])) |
| 173 | self.pos_drop = nn.Dropout(p=drop_rate) |
| 174 | |
| 175 | transformers = [] |
| 176 | # stochastic depth decay rule |
| 177 | dpr = [x.tolist() for x in torch.linspace(0, drop_path_rate, sum(depth)).split(depth)] |
| 178 | for stage in range(len(depth)): |
| 179 | pool = None |
| 180 | if stage < len(heads) - 1: |
| 181 | pool = ConvHeadPooling( |
| 182 | base_dims[stage] * heads[stage], base_dims[stage + 1] * heads[stage + 1], stride=2) |
| 183 | transformers += [Transformer( |
| 184 | base_dims[stage], depth[stage], heads[stage], mlp_ratio, pool=pool, |
| 185 | drop_rate=drop_rate, attn_drop_rate=attn_drop_rate, drop_path_prob=dpr[stage]) |
| 186 | ] |
| 187 | self.transformers = SequentialTuple(*transformers) |
| 188 | self.norm = nn.LayerNorm(base_dims[-1] * heads[-1], eps=1e-6) |
| 189 | self.num_features = self.embed_dim = base_dims[-1] * heads[-1] |
| 190 | |
| 191 | # Classifier head |
| 192 | self.head = nn.Linear(self.embed_dim, num_classes) if num_classes > 0 else nn.Identity() |
| 193 | self.head_dist = None |
| 194 | if distilled: |
| 195 | self.head_dist = nn.Linear(self.embed_dim, self.num_classes) if num_classes > 0 else nn.Identity() |
| 196 | |
| 197 | trunc_normal_(self.pos_embed, std=.02) |
| 198 | trunc_normal_(self.cls_token, std=.02) |
| 199 | self.apply(self._init_weights) |
| 200 | |
| 201 | def _init_weights(self, m): |
| 202 | if isinstance(m, nn.LayerNorm): |
nothing calls this directly
no test coverage detected