Method init_weights

mmseg/models/backbones/vit.py:289–336

(self)

        return getattr(self, self.norm1_name)

    def init_weights(self):
        if isinstance(self.init_cfg, dict) and self.init_cfg.get("type") == "Pretrained":
            logger = get_root_logger()
            checkpoint = CheckpointLoader.load_checkpoint(
                self.init_cfg["checkpoint"], logger=logger, map_location="cpu"
            )

            if "state_dict" in checkpoint:
                state_dict = checkpoint["state_dict"]
            else:
                state_dict = checkpoint

            if "pos_embed" in state_dict.keys():
                if self.pos_embed.shape != state_dict["pos_embed"].shape:
                    logger.info(
                        msg=f"Resize the pos_embed shape from "
                        f"{state_dict['pos_embed'].shape} to "
                        f"{self.pos_embed.shape}"
                    )
                    h, w = self.img_size
                    pos_size = int(math.sqrt(state_dict["pos_embed"].shape[1] - 1))
                    state_dict["pos_embed"] = self.resize_pos_embed(
                        state_dict["pos_embed"],
                        (h // self.patch_size, w // self.patch_size),
                        (pos_size, pos_size),
                        self.interpolate_mode,
                    )

            load_state_dict(self, state_dict, strict=False, logger=logger)
        elif self.init_cfg is not None:
            super(VisionTransformer, self).init_weights()
        else:
            # We only implement the 'jax_impl' initialization implemented at
            # https://github.com/rwightman/pytorch-image-models/blob/master/timm/models/vision_transformer.py#L353  # noqa: E501
            trunc_normal_(self.pos_embed, std=0.02)
            trunc_normal_(self.cls_token, std=0.02)
            for n, m in self.named_modules():
                if isinstance(m, nn.Linear):
                    trunc_normal_(m.weight, std=0.02)
                    if m.bias is not None:
                        if "ffn" in n:
                            nn.init.normal_(m.bias, mean=0.0, std=1e-6)
                        else:
                            nn.init.constant_(m.bias, 0)
                elif isinstance(m, nn.Conv2d):
                    kaiming_init(m, mode="fan_in", bias=0.0)
                elif isinstance(m, (_BatchNorm, nn.GroupNorm, nn.LayerNorm)):
                    constant_init(m, val=1.0, bias=0.0)

    def _pos_embeding(self, patched_img, hw_shape, pos_embed):
        """Positioning embeding method.