(self)
| 260 | return getattr(self, self.norm1_name) |
| 261 | |
| 262 | def init_weights(self): |
| 263 | if isinstance(self.pretrained, str): |
| 264 | logger = get_root_logger() |
| 265 | checkpoint = _load_checkpoint( |
| 266 | self.pretrained, logger=logger, map_location='cpu') |
| 267 | if 'state_dict' in checkpoint: |
| 268 | state_dict = checkpoint['state_dict'] |
| 269 | else: |
| 270 | state_dict = checkpoint |
| 271 | |
| 272 | if 'pos_embed' in state_dict.keys(): |
| 273 | if self.pos_embed.shape != state_dict['pos_embed'].shape: |
| 274 | logger.info(msg=f'Resize the pos_embed shape from ' |
| 275 | f'{state_dict["pos_embed"].shape} to ' |
| 276 | f'{self.pos_embed.shape}') |
| 277 | h, w = self.img_size |
| 278 | pos_size = int( |
| 279 | math.sqrt(state_dict['pos_embed'].shape[1])) |
| 280 | state_dict['pos_embed'] = self.resize_pos_embed( |
| 281 | state_dict['pos_embed'], |
| 282 | (h // self.patch_size, w // self.patch_size), |
| 283 | (pos_size, pos_size), self.interpolate_mode) |
| 284 | |
| 285 | self.load_state_dict(state_dict, False) |
| 286 | |
| 287 | elif self.pretrained is None: |
| 288 | super(VisionTransformer, self).init_weights() |
| 289 | # We only implement the 'jax_impl' initialization implemented at |
| 290 | # https://github.com/rwightman/pytorch-image-models/blob/master/timm/models/vision_transformer.py#L353 # noqa: E501 |
| 291 | trunc_normal_init(self.pos_embed, std=.02) |
| 292 | # trunc_normal_init(self.cls_token, std=.02) |
| 293 | for n, m in self.named_modules(): |
| 294 | if isinstance(m, nn.Linear): |
| 295 | trunc_normal_init(m.weight, std=.02) |
| 296 | if m.bias is not None: |
| 297 | if 'ffn' in n: |
| 298 | normal_init(m.bias, std=1e-6) |
| 299 | else: |
| 300 | constant_init(m.bias, 0) |
| 301 | elif isinstance(m, nn.Conv2d): |
| 302 | kaiming_init(m.weight, mode='fan_in') |
| 303 | if m.bias is not None: |
| 304 | constant_init(m.bias, 0) |
| 305 | elif isinstance(m, (_BatchNorm, nn.GroupNorm, nn.LayerNorm)): |
| 306 | constant_init(m.bias, 0) |
| 307 | constant_init(m.weight, 1.0) |
| 308 | |
| 309 | def _pos_embeding(self, patched_img, hw_shape, pos_embed): |
| 310 | """Positiong embeding method. |
nothing calls this directly
no test coverage detected