@auto_docstring

class Dinov2WithRegistersModel(Dinov2WithRegistersPreTrainedModel):

    def __init__(self, config: Dinov2WithRegistersConfig):

        super().__init__(config)

        self.config = config



        self.embeddings = Dinov2WithRegistersEmbeddings(config)

        self.encoder = Dinov2WithRegistersEncoder(config)



        self.layernorm = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps)



        # Initialize weights and apply final processing

        self.post_init()



    def get_input_embeddings(self) -> Dinov2WithRegistersPatchEmbeddings:

        return self.embeddings.patch_embeddings



    def _prune_heads(self, heads_to_prune: dict[int, list[int]]) -> None:

        """

        Prunes heads of the model. heads_to_prune: dict of {layer_num: list of heads to prune in this layer} See base

        class PreTrainedModel

        """

        for layer, heads in heads_to_prune.items():

            self.encoder.layer[layer].attention.prune_heads(heads)



    @auto_docstring

    def forward(

        self,

        pixel_values: Optional[torch.Tensor] = None,

        bool_masked_pos: Optional[torch.Tensor] = None,

        head_mask: Optional[torch.Tensor] = None,

        output_attentions: Optional[bool] = None,

        output_hidden_states: Optional[bool] = None,

        return_dict: Optional[bool] = None,

    ) -> Union[tuple, BaseModelOutputWithPooling]:

        r"""

        bool_masked_pos (`torch.BoolTensor` of shape `(batch_size, sequence_length)`):

            Boolean masked positions. Indicates which patches are masked (1) and which aren't (0). Only relevant for

            pre-training.

        """

        output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions

        output_hidden_states = (

            output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states

        )

        return_dict = return_dict if return_dict is not None else self.config.use_return_dict



        if pixel_values is None:

            raise ValueError("You have to specify pixel_values")



        # Prepare head mask if needed

        # 1.0 in head_mask indicate we keep the head

        # attention_probs has shape bsz x n_heads x N x N

        # input head_mask has shape [num_heads] or [num_hidden_layers x num_heads]

        # and head_mask is converted to shape [num_hidden_layers x batch x num_heads x seq_length x seq_length]

        head_mask = self.get_head_mask(head_mask, self.config.num_hidden_layers)



        embedding_output = self.embeddings(pixel_values, bool_masked_pos=bool_masked_pos)



        encoder_outputs = self.encoder(

            embedding_output,

            head_mask=head_mask,

            output_attentions=output_attentions,

            output_hidden_states=output_hidden_states,

            return_dict=return_dict,

        )

        sequence_output = encoder_outputs[0]

        sequence_output = self.layernorm(sequence_output)

        pooled_output = sequence_output[:, 0, :]



        if not return_dict:

            head_outputs = (sequence_output, pooled_output)

            return head_outputs + encoder_outputs[1:]



        return BaseModelOutputWithPooling(

            last_hidden_state=sequence_output,

            pooler_output=pooled_output,

            hidden_states=encoder_outputs.hidden_states,

            attentions=encoder_outputs.attentions,

        )



再来结合pytorch实现判断一下（我之前使用的）:



class DinoVisionTransformer(nn.Module):

    def __init__(

        self,

        img_size=224,

        patch_size=16,

        in_chans=3,

        embed_dim=768,

        depth=12,

        num_heads=12,

        mlp_ratio=4.0,

        qkv_bias=True,

        ffn_bias=True,

        proj_bias=True,

        drop_path_rate=0.0,

        drop_path_uniform=False,

        init_values=None,  # for layerscale: None or 0 => no layerscale

        embed_layer=PatchEmbed,

        act_layer=nn.GELU,

        block_fn=Block,

        ffn_layer="mlp",

        block_chunks=1,

        num_register_tokens=0,

        interpolate_antialias=False,

        interpolate_offset=0.1,

    ):

        """

        Args:

            img_size (int, tuple): input image size

            patch_size (int, tuple): patch size

            in_chans (int): number of input channels

            embed_dim (int): embedding dimension

            depth (int): depth of transformer

            num_heads (int): number of attention heads

            mlp_ratio (int): ratio of mlp hidden dim to embedding dim

            qkv_bias (bool): enable bias for qkv if True

            proj_bias (bool): enable bias for proj in attn if True

            ffn_bias (bool): enable bias for ffn if True

            drop_path_rate (float): stochastic depth rate

            drop_path_uniform (bool): apply uniform drop rate across blocks

            weight_init (str): weight init scheme

            init_values (float): layer-scale init values

            embed_layer (nn.Module): patch embedding layer

            act_layer (nn.Module): MLP activation layer

            block_fn (nn.Module): transformer block class

            ffn_layer (str): "mlp", "swiglu", "swiglufused" or "identity"

            block_chunks: (int) split block sequence into block_chunks units for FSDP wrap

            num_register_tokens: (int) number of extra cls tokens (so-called "registers")

            interpolate_antialias: (str) flag to apply anti-aliasing when interpolating positional embeddings

            interpolate_offset: (float) work-around offset to apply when interpolating positional embeddings

        """

        super().__init__()

        norm_layer = partial(nn.LayerNorm, eps=1e-6)



        self.num_features = self.embed_dim = embed_dim  # num_features for consistency with other models

        self.num_tokens = 1

        self.n_blocks = depth

        self.num_heads = num_heads

        self.patch_size = patch_size

        self.num_register_tokens = num_register_tokens

        self.interpolate_antialias = interpolate_antialias

        self.interpolate_offset = interpolate_offset



        self.patch_embed = embed_layer(img_size=img_size, patch_size=patch_size, in_chans=in_chans, embed_dim=embed_dim)

        num_patches = self.patch_embed.num_patches



        self.cls_token = nn.Parameter(torch.zeros(1, 1, embed_dim))

        self.pos_embed = nn.Parameter(torch.zeros(1, num_patches + self.num_tokens, embed_dim))

        assert num_register_tokens >= 0

        self.register_tokens = (

            nn.Parameter(torch.zeros(1, num_register_tokens, embed_dim)) if num_register_tokens else None

        )



        if drop_path_uniform is True:

            dpr = [drop_path_rate] * depth

        else:

            dpr = np.linspace(0, drop_path_rate, depth).tolist()  # stochastic depth decay rule



        if ffn_layer == "mlp":

            logger.info("using MLP layer as FFN")

            ffn_layer = Mlp

        elif ffn_layer == "swiglufused" or ffn_layer == "swiglu":

            logger.info("using SwiGLU layer as FFN")

            ffn_layer = SwiGLUFFNFused

        elif ffn_layer == "identity":

            logger.info("using Identity layer as FFN")



            def f(*args, **kwargs):

                return nn.Identity()



            ffn_layer = f

        else:

            raise NotImplementedError



        blocks_list = [

            block_fn(

                dim=embed_dim,

                num_heads=num_heads,

                mlp_ratio=mlp_ratio,

                qkv_bias=qkv_bias,

                proj_bias=proj_bias,

                ffn_bias=ffn_bias,

                drop_path=dpr[i],

                norm_layer=norm_layer,

                act_layer=act_layer,

                ffn_layer=ffn_layer,

                init_values=init_values,

            )

            for i in range(depth)

        ]

        if block_chunks > 0:

            self.chunked_blocks = True

            chunked_blocks = []

            chunksize = depth // block_chunks

            for i in range(0, depth, chunksize):

                # this is to keep the block index consistent if we chunk the block list

                chunked_blocks.append([nn.Identity()] * i + blocks_list[i : i + chunksize])

            self.blocks = nn.ModuleList([BlockChunk(p) for p in chunked_blocks])

        else:

            self.chunked_blocks = False

            self.blocks = nn.ModuleList(blocks_list)



        self.norm = norm_layer(embed_dim)

        self.head = nn.Identity()



        self.mask_token = nn.Parameter(torch.zeros(1, embed_dim))



        self.init_weights()



    def init_weights(self):

        trunc_normal_(self.pos_embed, std=0.02)

        nn.init.normal_(self.cls_token, std=1e-6)

        if self.register_tokens is not None:

            nn.init.normal_(self.register_tokens, std=1e-6)

        named_apply(init_weights_vit_timm, self)



    def interpolate_pos_encoding(self, x, w, h):

        previous_dtype = x.dtype

        npatch = x.shape[1] - 1

        N = self.pos_embed.shape[1] - 1

        if npatch == N and w == h:

            return self.pos_embed

        pos_embed = self.pos_embed.float()

        class_pos_embed = pos_embed[:, 0]

        patch_pos_embed = pos_embed[:, 1:]

        dim = x.shape[-1]

        w0 = w // self.patch_size

        h0 = h // self.patch_size

        M = int(math.sqrt(N))  # Recover the number of patches in each dimension

        assert N == M * M

        kwargs = {}

        if self.interpolate_offset:

            # Historical kludge: add a small number to avoid floating point error in the interpolation, see https://github.com/facebookresearch/dino/issues/8

            # Note: still needed for backward-compatibility, the underlying operators are using both output size and scale factors

            sx = float(w0 + self.interpolate_offset) / M

            sy = float(h0 + self.interpolate_offset) / M

            kwargs["scale_factor"] = (sx, sy)

        else:

            # Simply specify an output size instead of a scale factor

            kwargs["size"] = (w0, h0)

        patch_pos_embed = nn.functional.interpolate(

            patch_pos_embed.reshape(1, M, M, dim).permute(0, 3, 1, 2),

            mode="bicubic",

            antialias=self.interpolate_antialias,

            **kwargs,

        )

        assert (w0, h0) == patch_pos_embed.shape[-2:]

        patch_pos_embed = patch_pos_embed.permute(0, 2, 3, 1).view(1, -1, dim)

        return torch.cat((class_pos_embed.unsqueeze(0), patch_pos_embed), dim=1).to(previous_dtype)



    def prepare_tokens_with_masks(self, x, masks=None):

        B, nc, w, h = x.shape

        x = self.patch_embed(x)

        if masks is not None:

            x = torch.where(masks.unsqueeze(-1), self.mask_token.to(x.dtype).unsqueeze(0), x)



        x = torch.cat((self.cls_token.expand(x.shape[0], -1, -1), x), dim=1)

        x = x + self.interpolate_pos_encoding(x, w, h)



        if self.register_tokens is not None:

            x = torch.cat(

                (

                    x[:, :1],

                    self.register_tokens.expand(x.shape[0], -1, -1),

                    x[:, 1:],

                ),

                dim=1,

            )



        return x



    def forward_features_list(self, x_list, masks_list):

        x = [self.prepare_tokens_with_masks(x, masks) for x, masks in zip(x_list, masks_list)]

        for blk in self.blocks:

            x = blk(x)



        all_x = x

        output = []

        for x, masks in zip(all_x, masks_list):

            x_norm = self.norm(x)

            output.append(

                {

                    "x_norm_clstoken": x_norm[:, 0],

                    "x_norm_regtokens": x_norm[:, 1 : self.num_register_tokens + 1],

                    "x_norm_patchtokens": x_norm[:, self.num_register_tokens + 1 :],

                    "x_prenorm": x,

                    "masks": masks,

                }

            )

        return output



    def forward_features(self, x, masks=None):

        if isinstance(x, list):

            return self.forward_features_list(x, masks)



        x = self.prepare_tokens_with_masks(x, masks)



        for blk in self.blocks:

            x = blk(x)



        x_norm = self.norm(x)

        return {

            "x_norm_clstoken": x_norm[:, 0],

            "x_norm_regtokens": x_norm[:, 1 : self.num_register_tokens + 1],

            "x_norm_patchtokens": x_norm[:, self.num_register_tokens + 1 :],

            "x_prenorm": x,

            "masks": masks,

        }



并回答两个问题：

为什么outputs.pooler_output是cls

我如果要在替换后（hf实现）的实现中提取x_norm_regtokens和x_norm_patchtokens，顺序与"x_norm_regtokens": x_norm[:, 1 : self.num_register_tokens + 1], "x_norm_patchtokens": x_norm[:, self.num_register_tokens + 1 :]一致吗