Aluode/PerceptionLabPortable

1# coding=utf-82# Copyright 2022 Microsoft Research and The HuggingFace Inc. team. All rights reserved.3#4# Licensed under the Apache License, Version 2.0 (the "License");5# you may not use this file except in compliance with the License.6# You may obtain a copy of the License at7#8#     http://www.apache.org/licenses/LICENSE-2.09#10# Unless required by applicable law or agreed to in writing, software11# distributed under the License is distributed on an "AS IS" BASIS,12# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.13# See the License for the specific language governing permissions and14# limitations under the License.15"""PyTorch CvT model."""16 17import collections.abc18from dataclasses import dataclass19from typing import Optional, Union20 21import torch22from torch import nn23from torch.nn import BCEWithLogitsLoss, CrossEntropyLoss, MSELoss24 25from ...modeling_outputs import ImageClassifierOutputWithNoAttention, ModelOutput26from ...modeling_utils import PreTrainedModel27from ...pytorch_utils import find_pruneable_heads_and_indices, prune_linear_layer28from ...utils import auto_docstring, logging29from .configuration_cvt import CvtConfig30 31 32logger = logging.get_logger(__name__)33 34 35@dataclass36@auto_docstring(37    custom_intro="""38    Base class for model's outputs, with potential hidden states and attentions.39    """40)41class BaseModelOutputWithCLSToken(ModelOutput):42    r"""43    cls_token_value (`torch.FloatTensor` of shape `(batch_size, 1, hidden_size)`):44        Classification token at the output of the last layer of the model.45    """46 47    last_hidden_state: Optional[torch.FloatTensor] = None48    cls_token_value: Optional[torch.FloatTensor] = None49    hidden_states: Optional[tuple[torch.FloatTensor, ...]] = None50 51 52# Copied from transformers.models.beit.modeling_beit.drop_path53def drop_path(input: torch.Tensor, drop_prob: float = 0.0, training: bool = False) -> torch.Tensor:54    """55    Drop paths (Stochastic Depth) per sample (when applied in main path of residual blocks).56 57    Comment by Ross Wightman: This is the same as the DropConnect impl I created for EfficientNet, etc networks,58    however, the original name is misleading as 'Drop Connect' is a different form of dropout in a separate paper...59    See discussion: https://github.com/tensorflow/tpu/issues/494#issuecomment-532968956 ... I've opted for changing the60    layer and argument names to 'drop path' rather than mix DropConnect as a layer name and use 'survival rate' as the61    argument.62    """63    if drop_prob == 0.0 or not training:64        return input65    keep_prob = 1 - drop_prob66    shape = (input.shape[0],) + (1,) * (input.ndim - 1)  # work with diff dim tensors, not just 2D ConvNets67    random_tensor = keep_prob + torch.rand(shape, dtype=input.dtype, device=input.device)68    random_tensor.floor_()  # binarize69    output = input.div(keep_prob) * random_tensor70    return output71 72 73# Copied from transformers.models.beit.modeling_beit.BeitDropPath74class CvtDropPath(nn.Module):75    """Drop paths (Stochastic Depth) per sample (when applied in main path of residual blocks)."""76 77    def __init__(self, drop_prob: Optional[float] = None) -> None:78        super().__init__()79        self.drop_prob = drop_prob80 81    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:82        return drop_path(hidden_states, self.drop_prob, self.training)83 84    def extra_repr(self) -> str:85        return f"p={self.drop_prob}"86 87 88class CvtEmbeddings(nn.Module):89    """90    Construct the CvT embeddings.91    """92 93    def __init__(self, patch_size, num_channels, embed_dim, stride, padding, dropout_rate):94        super().__init__()95        self.convolution_embeddings = CvtConvEmbeddings(96            patch_size=patch_size, num_channels=num_channels, embed_dim=embed_dim, stride=stride, padding=padding97        )98        self.dropout = nn.Dropout(dropout_rate)99 100    def forward(self, pixel_values):101        hidden_state = self.convolution_embeddings(pixel_values)102        hidden_state = self.dropout(hidden_state)103        return hidden_state104 105 106class CvtConvEmbeddings(nn.Module):107    """108    Image to Conv Embedding.109    """110 111    def __init__(self, patch_size, num_channels, embed_dim, stride, padding):112        super().__init__()113        patch_size = patch_size if isinstance(patch_size, collections.abc.Iterable) else (patch_size, patch_size)114        self.patch_size = patch_size115        self.projection = nn.Conv2d(num_channels, embed_dim, kernel_size=patch_size, stride=stride, padding=padding)116        self.normalization = nn.LayerNorm(embed_dim)117 118    def forward(self, pixel_values):119        pixel_values = self.projection(pixel_values)120        batch_size, num_channels, height, width = pixel_values.shape121        hidden_size = height * width122        # rearrange "b c h w -> b (h w) c"123        pixel_values = pixel_values.view(batch_size, num_channels, hidden_size).permute(0, 2, 1)124        if self.normalization:125            pixel_values = self.normalization(pixel_values)126        # rearrange "b (h w) c" -> b c h w"127        pixel_values = pixel_values.permute(0, 2, 1).view(batch_size, num_channels, height, width)128        return pixel_values129 130 131class CvtSelfAttentionConvProjection(nn.Module):132    def __init__(self, embed_dim, kernel_size, padding, stride):133        super().__init__()134        self.convolution = nn.Conv2d(135            embed_dim,136            embed_dim,137            kernel_size=kernel_size,138            padding=padding,139            stride=stride,140            bias=False,141            groups=embed_dim,142        )143        self.normalization = nn.BatchNorm2d(embed_dim)144 145    def forward(self, hidden_state):146        hidden_state = self.convolution(hidden_state)147        hidden_state = self.normalization(hidden_state)148        return hidden_state149 150 151class CvtSelfAttentionLinearProjection(nn.Module):152    def forward(self, hidden_state):153        batch_size, num_channels, height, width = hidden_state.shape154        hidden_size = height * width155        # rearrange " b c h w -> b (h w) c"156        hidden_state = hidden_state.view(batch_size, num_channels, hidden_size).permute(0, 2, 1)157        return hidden_state158 159 160class CvtSelfAttentionProjection(nn.Module):161    def __init__(self, embed_dim, kernel_size, padding, stride, projection_method="dw_bn"):162        super().__init__()163        if projection_method == "dw_bn":164            self.convolution_projection = CvtSelfAttentionConvProjection(embed_dim, kernel_size, padding, stride)165        self.linear_projection = CvtSelfAttentionLinearProjection()166 167    def forward(self, hidden_state):168        hidden_state = self.convolution_projection(hidden_state)169        hidden_state = self.linear_projection(hidden_state)170        return hidden_state171 172 173class CvtSelfAttention(nn.Module):174    def __init__(175        self,176        num_heads,177        embed_dim,178        kernel_size,179        padding_q,180        padding_kv,181        stride_q,182        stride_kv,183        qkv_projection_method,184        qkv_bias,185        attention_drop_rate,186        with_cls_token=True,187        **kwargs,188    ):189        super().__init__()190        self.scale = embed_dim**-0.5191        self.with_cls_token = with_cls_token192        self.embed_dim = embed_dim193        self.num_heads = num_heads194 195        self.convolution_projection_query = CvtSelfAttentionProjection(196            embed_dim,197            kernel_size,198            padding_q,199            stride_q,200            projection_method="linear" if qkv_projection_method == "avg" else qkv_projection_method,201        )202        self.convolution_projection_key = CvtSelfAttentionProjection(203            embed_dim, kernel_size, padding_kv, stride_kv, projection_method=qkv_projection_method204        )205        self.convolution_projection_value = CvtSelfAttentionProjection(206            embed_dim, kernel_size, padding_kv, stride_kv, projection_method=qkv_projection_method207        )208 209        self.projection_query = nn.Linear(embed_dim, embed_dim, bias=qkv_bias)210        self.projection_key = nn.Linear(embed_dim, embed_dim, bias=qkv_bias)211        self.projection_value = nn.Linear(embed_dim, embed_dim, bias=qkv_bias)212 213        self.dropout = nn.Dropout(attention_drop_rate)214 215    def rearrange_for_multi_head_attention(self, hidden_state):216        batch_size, hidden_size, _ = hidden_state.shape217        head_dim = self.embed_dim // self.num_heads218        # rearrange 'b t (h d) -> b h t d'219        return hidden_state.view(batch_size, hidden_size, self.num_heads, head_dim).permute(0, 2, 1, 3)220 221    def forward(self, hidden_state, height, width):222        if self.with_cls_token:223            cls_token, hidden_state = torch.split(hidden_state, [1, height * width], 1)224        batch_size, hidden_size, num_channels = hidden_state.shape225        # rearrange "b (h w) c -> b c h w"226        hidden_state = hidden_state.permute(0, 2, 1).view(batch_size, num_channels, height, width)227 228        key = self.convolution_projection_key(hidden_state)229        query = self.convolution_projection_query(hidden_state)230        value = self.convolution_projection_value(hidden_state)231 232        if self.with_cls_token:233            query = torch.cat((cls_token, query), dim=1)234            key = torch.cat((cls_token, key), dim=1)235            value = torch.cat((cls_token, value), dim=1)236 237        head_dim = self.embed_dim // self.num_heads238 239        query = self.rearrange_for_multi_head_attention(self.projection_query(query))240        key = self.rearrange_for_multi_head_attention(self.projection_key(key))241        value = self.rearrange_for_multi_head_attention(self.projection_value(value))242 243        attention_score = torch.einsum("bhlk,bhtk->bhlt", [query, key]) * self.scale244        attention_probs = torch.nn.functional.softmax(attention_score, dim=-1)245        attention_probs = self.dropout(attention_probs)246 247        context = torch.einsum("bhlt,bhtv->bhlv", [attention_probs, value])248        # rearrange"b h t d -> b t (h d)"249        _, _, hidden_size, _ = context.shape250        context = context.permute(0, 2, 1, 3).contiguous().view(batch_size, hidden_size, self.num_heads * head_dim)251        return context252 253 254class CvtSelfOutput(nn.Module):255    """256    The residual connection is defined in CvtLayer instead of here (as is the case with other models), due to the257    layernorm applied before each block.258    """259 260    def __init__(self, embed_dim, drop_rate):261        super().__init__()262        self.dense = nn.Linear(embed_dim, embed_dim)263        self.dropout = nn.Dropout(drop_rate)264 265    def forward(self, hidden_state, input_tensor):266        hidden_state = self.dense(hidden_state)267        hidden_state = self.dropout(hidden_state)268        return hidden_state269 270 271class CvtAttention(nn.Module):272    def __init__(273        self,274        num_heads,275        embed_dim,276        kernel_size,277        padding_q,278        padding_kv,279        stride_q,280        stride_kv,281        qkv_projection_method,282        qkv_bias,283        attention_drop_rate,284        drop_rate,285        with_cls_token=True,286    ):287        super().__init__()288        self.attention = CvtSelfAttention(289            num_heads,290            embed_dim,291            kernel_size,292            padding_q,293            padding_kv,294            stride_q,295            stride_kv,296            qkv_projection_method,297            qkv_bias,298            attention_drop_rate,299            with_cls_token,300        )301        self.output = CvtSelfOutput(embed_dim, drop_rate)302        self.pruned_heads = set()303 304    def prune_heads(self, heads):305        if len(heads) == 0:306            return307        heads, index = find_pruneable_heads_and_indices(308            heads, self.attention.num_attention_heads, self.attention.attention_head_size, self.pruned_heads309        )310 311        # Prune linear layers312        self.attention.query = prune_linear_layer(self.attention.query, index)313        self.attention.key = prune_linear_layer(self.attention.key, index)314        self.attention.value = prune_linear_layer(self.attention.value, index)315        self.output.dense = prune_linear_layer(self.output.dense, index, dim=1)316 317        # Update hyper params and store pruned heads318        self.attention.num_attention_heads = self.attention.num_attention_heads - len(heads)319        self.attention.all_head_size = self.attention.attention_head_size * self.attention.num_attention_heads320        self.pruned_heads = self.pruned_heads.union(heads)321 322    def forward(self, hidden_state, height, width):323        self_output = self.attention(hidden_state, height, width)324        attention_output = self.output(self_output, hidden_state)325        return attention_output326 327 328class CvtIntermediate(nn.Module):329    def __init__(self, embed_dim, mlp_ratio):330        super().__init__()331        self.dense = nn.Linear(embed_dim, int(embed_dim * mlp_ratio))332        self.activation = nn.GELU()333 334    def forward(self, hidden_state):335        hidden_state = self.dense(hidden_state)336        hidden_state = self.activation(hidden_state)337        return hidden_state338 339 340class CvtOutput(nn.Module):341    def __init__(self, embed_dim, mlp_ratio, drop_rate):342        super().__init__()343        self.dense = nn.Linear(int(embed_dim * mlp_ratio), embed_dim)344        self.dropout = nn.Dropout(drop_rate)345 346    def forward(self, hidden_state, input_tensor):347        hidden_state = self.dense(hidden_state)348        hidden_state = self.dropout(hidden_state)349        hidden_state = hidden_state + input_tensor350        return hidden_state351 352 353class CvtLayer(nn.Module):354    """355    CvtLayer composed by attention layers, normalization and multi-layer perceptrons (mlps).356    """357 358    def __init__(359        self,360        num_heads,361        embed_dim,362        kernel_size,363        padding_q,364        padding_kv,365        stride_q,366        stride_kv,367        qkv_projection_method,368        qkv_bias,369        attention_drop_rate,370        drop_rate,371        mlp_ratio,372        drop_path_rate,373        with_cls_token=True,374    ):375        super().__init__()376        self.attention = CvtAttention(377            num_heads,378            embed_dim,379            kernel_size,380            padding_q,381            padding_kv,382            stride_q,383            stride_kv,384            qkv_projection_method,385            qkv_bias,386            attention_drop_rate,387            drop_rate,388            with_cls_token,389        )390 391        self.intermediate = CvtIntermediate(embed_dim, mlp_ratio)392        self.output = CvtOutput(embed_dim, mlp_ratio, drop_rate)393        self.drop_path = CvtDropPath(drop_prob=drop_path_rate) if drop_path_rate > 0.0 else nn.Identity()394        self.layernorm_before = nn.LayerNorm(embed_dim)395        self.layernorm_after = nn.LayerNorm(embed_dim)396 397    def forward(self, hidden_state, height, width):398        self_attention_output = self.attention(399            self.layernorm_before(hidden_state),  # in Cvt, layernorm is applied before self-attention400            height,401            width,402        )403        attention_output = self_attention_output404        attention_output = self.drop_path(attention_output)405 406        # first residual connection407        hidden_state = attention_output + hidden_state408 409        # in Cvt, layernorm is also applied after self-attention410        layer_output = self.layernorm_after(hidden_state)411        layer_output = self.intermediate(layer_output)412 413        # second residual connection is done here414        layer_output = self.output(layer_output, hidden_state)415        layer_output = self.drop_path(layer_output)416        return layer_output417 418 419class CvtStage(nn.Module):420    def __init__(self, config, stage):421        super().__init__()422        self.config = config423        self.stage = stage424        if self.config.cls_token[self.stage]:425            self.cls_token = nn.Parameter(torch.randn(1, 1, self.config.embed_dim[-1]))426 427        self.embedding = CvtEmbeddings(428            patch_size=config.patch_sizes[self.stage],429            stride=config.patch_stride[self.stage],430            num_channels=config.num_channels if self.stage == 0 else config.embed_dim[self.stage - 1],431            embed_dim=config.embed_dim[self.stage],432            padding=config.patch_padding[self.stage],433            dropout_rate=config.drop_rate[self.stage],434        )435 436        drop_path_rates = [437            x.item() for x in torch.linspace(0, config.drop_path_rate[self.stage], config.depth[stage], device="cpu")438        ]439 440        self.layers = nn.Sequential(441            *[442                CvtLayer(443                    num_heads=config.num_heads[self.stage],444                    embed_dim=config.embed_dim[self.stage],445                    kernel_size=config.kernel_qkv[self.stage],446                    padding_q=config.padding_q[self.stage],447                    padding_kv=config.padding_kv[self.stage],448                    stride_kv=config.stride_kv[self.stage],449                    stride_q=config.stride_q[self.stage],450                    qkv_projection_method=config.qkv_projection_method[self.stage],451                    qkv_bias=config.qkv_bias[self.stage],452                    attention_drop_rate=config.attention_drop_rate[self.stage],453                    drop_rate=config.drop_rate[self.stage],454                    drop_path_rate=drop_path_rates[self.stage],455                    mlp_ratio=config.mlp_ratio[self.stage],456                    with_cls_token=config.cls_token[self.stage],457                )458                for _ in range(config.depth[self.stage])459            ]460        )461 462    def forward(self, hidden_state):463        cls_token = None464        hidden_state = self.embedding(hidden_state)465        batch_size, num_channels, height, width = hidden_state.shape466        # rearrange b c h w -> b (h w) c"467        hidden_state = hidden_state.view(batch_size, num_channels, height * width).permute(0, 2, 1)468        if self.config.cls_token[self.stage]:469            cls_token = self.cls_token.expand(batch_size, -1, -1)470            hidden_state = torch.cat((cls_token, hidden_state), dim=1)471 472        for layer in self.layers:473            layer_outputs = layer(hidden_state, height, width)474            hidden_state = layer_outputs475 476        if self.config.cls_token[self.stage]:477            cls_token, hidden_state = torch.split(hidden_state, [1, height * width], 1)478        hidden_state = hidden_state.permute(0, 2, 1).view(batch_size, num_channels, height, width)479        return hidden_state, cls_token480 481 482class CvtEncoder(nn.Module):483    def __init__(self, config):484        super().__init__()485        self.config = config486        self.stages = nn.ModuleList([])487        for stage_idx in range(len(config.depth)):488            self.stages.append(CvtStage(config, stage_idx))489 490    def forward(self, pixel_values, output_hidden_states=False, return_dict=True):491        all_hidden_states = () if output_hidden_states else None492        hidden_state = pixel_values493 494        cls_token = None495        for _, (stage_module) in enumerate(self.stages):496            hidden_state, cls_token = stage_module(hidden_state)497            if output_hidden_states:498                all_hidden_states = all_hidden_states + (hidden_state,)499 500        if not return_dict:501            return tuple(v for v in [hidden_state, cls_token, all_hidden_states] if v is not None)502 503        return BaseModelOutputWithCLSToken(504            last_hidden_state=hidden_state,505            cls_token_value=cls_token,506            hidden_states=all_hidden_states,507        )508 509 510@auto_docstring511class CvtPreTrainedModel(PreTrainedModel):512    config: CvtConfig513    base_model_prefix = "cvt"514    main_input_name = "pixel_values"515    _no_split_modules = ["CvtLayer"]516 517    def _init_weights(self, module):518        """Initialize the weights"""519        if isinstance(module, (nn.Linear, nn.Conv2d)):520            module.weight.data = nn.init.trunc_normal_(module.weight.data, mean=0.0, std=self.config.initializer_range)521            if module.bias is not None:522                module.bias.data.zero_()523        elif isinstance(module, nn.LayerNorm):524            module.bias.data.zero_()525            module.weight.data.fill_(1.0)526        elif isinstance(module, CvtStage):527            if self.config.cls_token[module.stage]:528                module.cls_token.data = nn.init.trunc_normal_(529                    module.cls_token.data, mean=0.0, std=self.config.initializer_range530                )531 532 533@auto_docstring534class CvtModel(CvtPreTrainedModel):535    def __init__(self, config, add_pooling_layer=True):536        r"""537        add_pooling_layer (bool, *optional*, defaults to `True`):538            Whether to add a pooling layer539        """540        super().__init__(config)541        self.config = config542        self.encoder = CvtEncoder(config)543        self.post_init()544 545    def _prune_heads(self, heads_to_prune):546        """547        Prunes heads of the model. heads_to_prune: dict of {layer_num: list of heads to prune in this layer} See base548        class PreTrainedModel549        """550        for layer, heads in heads_to_prune.items():551            self.encoder.layer[layer].attention.prune_heads(heads)552 553    @auto_docstring554    def forward(555        self,556        pixel_values: Optional[torch.Tensor] = None,557        output_hidden_states: Optional[bool] = None,558        return_dict: Optional[bool] = None,559    ) -> Union[tuple, BaseModelOutputWithCLSToken]:560        output_hidden_states = (561            output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states562        )563        return_dict = return_dict if return_dict is not None else self.config.use_return_dict564 565        if pixel_values is None:566            raise ValueError("You have to specify pixel_values")567 568        encoder_outputs = self.encoder(569            pixel_values,570            output_hidden_states=output_hidden_states,571            return_dict=return_dict,572        )573        sequence_output = encoder_outputs[0]574 575        if not return_dict:576            return (sequence_output,) + encoder_outputs[1:]577 578        return BaseModelOutputWithCLSToken(579            last_hidden_state=sequence_output,580            cls_token_value=encoder_outputs.cls_token_value,581            hidden_states=encoder_outputs.hidden_states,582        )583 584 585@auto_docstring(586    custom_intro="""587    Cvt Model transformer with an image classification head on top (a linear layer on top of the final hidden state of588    the [CLS] token) e.g. for ImageNet.589    """590)591class CvtForImageClassification(CvtPreTrainedModel):592    def __init__(self, config):593        super().__init__(config)594 595        self.num_labels = config.num_labels596        self.cvt = CvtModel(config, add_pooling_layer=False)597        self.layernorm = nn.LayerNorm(config.embed_dim[-1])598        # Classifier head599        self.classifier = (600            nn.Linear(config.embed_dim[-1], config.num_labels) if config.num_labels > 0 else nn.Identity()601        )602 603        # Initialize weights and apply final processing604        self.post_init()605 606    @auto_docstring607    def forward(608        self,609        pixel_values: Optional[torch.Tensor] = None,610        labels: Optional[torch.Tensor] = None,611        output_hidden_states: Optional[bool] = None,612        return_dict: Optional[bool] = None,613    ) -> Union[tuple, ImageClassifierOutputWithNoAttention]:614        r"""615        labels (`torch.LongTensor` of shape `(batch_size,)`, *optional*):616            Labels for computing the image classification/regression loss. Indices should be in `[0, ...,617            config.num_labels - 1]`. If `config.num_labels == 1` a regression loss is computed (Mean-Square loss), If618            `config.num_labels > 1` a classification loss is computed (Cross-Entropy).619        """620        return_dict = return_dict if return_dict is not None else self.config.use_return_dict621        outputs = self.cvt(622            pixel_values,623            output_hidden_states=output_hidden_states,624            return_dict=return_dict,625        )626 627        sequence_output = outputs[0]628        cls_token = outputs[1]629        if self.config.cls_token[-1]:630            sequence_output = self.layernorm(cls_token)631        else:632            batch_size, num_channels, height, width = sequence_output.shape633            # rearrange "b c h w -> b (h w) c"634            sequence_output = sequence_output.view(batch_size, num_channels, height * width).permute(0, 2, 1)635            sequence_output = self.layernorm(sequence_output)636 637        sequence_output_mean = sequence_output.mean(dim=1)638        logits = self.classifier(sequence_output_mean)639 640        loss = None641        if labels is not None:642            if self.config.problem_type is None:643                if self.config.num_labels == 1:644                    self.config.problem_type = "regression"645                elif self.config.num_labels > 1 and (labels.dtype == torch.long or labels.dtype == torch.int):646                    self.config.problem_type = "single_label_classification"647                else:648                    self.config.problem_type = "multi_label_classification"649 650            if self.config.problem_type == "regression":651                loss_fct = MSELoss()652                if self.config.num_labels == 1:653                    loss = loss_fct(logits.squeeze(), labels.squeeze())654                else:655                    loss = loss_fct(logits, labels)656            elif self.config.problem_type == "single_label_classification":657                loss_fct = CrossEntropyLoss()658                loss = loss_fct(logits.view(-1, self.config.num_labels), labels.view(-1))659            elif self.config.problem_type == "multi_label_classification":660                loss_fct = BCEWithLogitsLoss()661                loss = loss_fct(logits, labels)662 663        if not return_dict:664            output = (logits,) + outputs[2:]665            return ((loss,) + output) if loss is not None else output666 667        return ImageClassifierOutputWithNoAttention(loss=loss, logits=logits, hidden_states=outputs.hidden_states)668 669 670__all__ = ["CvtForImageClassification", "CvtModel", "CvtPreTrainedModel"]671