Aluode/PerceptionLabPortable

1# coding=utf-82# Copyright 2022 Sea AI Lab 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 PoolFormer model."""16 17import collections.abc18from typing import Optional, Union19 20import torch21from torch import nn22 23from ...activations import ACT2FN24from ...modeling_outputs import BaseModelOutputWithNoAttention, ImageClassifierOutputWithNoAttention25from ...modeling_utils import PreTrainedModel26from ...utils import auto_docstring, logging27from .configuration_poolformer import PoolFormerConfig28 29 30logger = logging.get_logger(__name__)31 32 33# Copied from transformers.models.beit.modeling_beit.drop_path34def drop_path(input: torch.Tensor, drop_prob: float = 0.0, training: bool = False) -> torch.Tensor:35    """36    Drop paths (Stochastic Depth) per sample (when applied in main path of residual blocks).37 38    Comment by Ross Wightman: This is the same as the DropConnect impl I created for EfficientNet, etc networks,39    however, the original name is misleading as 'Drop Connect' is a different form of dropout in a separate paper...40    See discussion: https://github.com/tensorflow/tpu/issues/494#issuecomment-532968956 ... I've opted for changing the41    layer and argument names to 'drop path' rather than mix DropConnect as a layer name and use 'survival rate' as the42    argument.43    """44    if drop_prob == 0.0 or not training:45        return input46    keep_prob = 1 - drop_prob47    shape = (input.shape[0],) + (1,) * (input.ndim - 1)  # work with diff dim tensors, not just 2D ConvNets48    random_tensor = keep_prob + torch.rand(shape, dtype=input.dtype, device=input.device)49    random_tensor.floor_()  # binarize50    output = input.div(keep_prob) * random_tensor51    return output52 53 54# Copied from transformers.models.beit.modeling_beit.BeitDropPath with Beit->PoolFormer55class PoolFormerDropPath(nn.Module):56    """Drop paths (Stochastic Depth) per sample (when applied in main path of residual blocks)."""57 58    def __init__(self, drop_prob: Optional[float] = None) -> None:59        super().__init__()60        self.drop_prob = drop_prob61 62    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:63        return drop_path(hidden_states, self.drop_prob, self.training)64 65    def extra_repr(self) -> str:66        return f"p={self.drop_prob}"67 68 69class PoolFormerEmbeddings(nn.Module):70    """71    Construct Patch Embeddings.72    """73 74    def __init__(self, hidden_size, num_channels, patch_size, stride, padding, norm_layer=None):75        super().__init__()76        patch_size = patch_size if isinstance(patch_size, collections.abc.Iterable) else (patch_size, patch_size)77        stride = stride if isinstance(stride, collections.abc.Iterable) else (stride, stride)78        padding = padding if isinstance(padding, collections.abc.Iterable) else (padding, padding)79 80        self.projection = nn.Conv2d(num_channels, hidden_size, kernel_size=patch_size, stride=stride, padding=padding)81        self.norm = norm_layer(hidden_size) if norm_layer else nn.Identity()82 83    def forward(self, pixel_values):84        embeddings = self.projection(pixel_values)85        embeddings = self.norm(embeddings)86        return embeddings87 88 89class PoolFormerGroupNorm(nn.GroupNorm):90    """91    Group Normalization with 1 group. Input: tensor in shape [B, C, H, W]92    """93 94    def __init__(self, num_channels, **kwargs):95        super().__init__(1, num_channels, **kwargs)96 97 98class PoolFormerPooling(nn.Module):99    def __init__(self, pool_size):100        super().__init__()101        self.pool = nn.AvgPool2d(pool_size, stride=1, padding=pool_size // 2, count_include_pad=False)102 103    def forward(self, hidden_states):104        return self.pool(hidden_states) - hidden_states105 106 107class PoolFormerOutput(nn.Module):108    def __init__(self, config, dropout_prob, hidden_size, intermediate_size):109        super().__init__()110        self.conv1 = nn.Conv2d(hidden_size, intermediate_size, 1)111        self.conv2 = nn.Conv2d(intermediate_size, hidden_size, 1)112        self.drop = PoolFormerDropPath(dropout_prob)113        if isinstance(config.hidden_act, str):114            self.act_fn = ACT2FN[config.hidden_act]115        else:116            self.act_fn = config.hidden_act117 118    def forward(self, hidden_states):119        hidden_states = self.conv1(hidden_states)120        hidden_states = self.act_fn(hidden_states)121        hidden_states = self.drop(hidden_states)122        hidden_states = self.conv2(hidden_states)123        hidden_states = self.drop(hidden_states)124 125        return hidden_states126 127 128class PoolFormerLayer(nn.Module):129    """This corresponds to the 'PoolFormerBlock' class in the original implementation."""130 131    def __init__(self, config, num_channels, pool_size, hidden_size, intermediate_size, drop_path):132        super().__init__()133        self.pooling = PoolFormerPooling(pool_size)134        self.output = PoolFormerOutput(config, drop_path, hidden_size, intermediate_size)135        self.before_norm = PoolFormerGroupNorm(num_channels)136        self.after_norm = PoolFormerGroupNorm(num_channels)137 138        # Useful for training neural nets139        self.drop_path = PoolFormerDropPath(drop_path) if drop_path > 0.0 else nn.Identity()140        self.use_layer_scale = config.use_layer_scale141        if config.use_layer_scale:142            self.layer_scale_1 = nn.Parameter(143                config.layer_scale_init_value * torch.ones(num_channels), requires_grad=True144            )145            self.layer_scale_2 = nn.Parameter(146                config.layer_scale_init_value * torch.ones(num_channels), requires_grad=True147            )148 149    def forward(self, hidden_states):150        if self.use_layer_scale:151            pooling_output = self.pooling(self.before_norm(hidden_states))152            scaled_op = self.layer_scale_1.unsqueeze(-1).unsqueeze(-1) * pooling_output153            # First residual connection154            hidden_states = hidden_states + self.drop_path(scaled_op)155            outputs = ()156 157            layer_output = self.output(self.after_norm(hidden_states))158            scaled_op = self.layer_scale_2.unsqueeze(-1).unsqueeze(-1) * layer_output159            # Second residual connection160            output = hidden_states + self.drop_path(scaled_op)161 162            outputs = (output,) + outputs163            return outputs164 165        else:166            pooling_output = self.drop_path(self.pooling(self.before_norm(hidden_states)))167            # First residual connection168            hidden_states = pooling_output + hidden_states169            outputs = ()170 171            # Second residual connection inside the PoolFormerOutput block172            layer_output = self.drop_path(self.output(self.after_norm(hidden_states)))173            output = hidden_states + layer_output174 175            outputs = (output,) + outputs176            return outputs177 178 179class PoolFormerEncoder(nn.Module):180    def __init__(self, config):181        super().__init__()182        self.config = config183        # stochastic depth decay rule184        dpr = [x.item() for x in torch.linspace(0, config.drop_path_rate, sum(config.depths), device="cpu")]185 186        # patch embeddings187        embeddings = []188        for i in range(config.num_encoder_blocks):189            embeddings.append(190                PoolFormerEmbeddings(191                    patch_size=config.patch_sizes[i],192                    stride=config.strides[i],193                    padding=config.padding[i],194                    num_channels=config.num_channels if i == 0 else config.hidden_sizes[i - 1],195                    hidden_size=config.hidden_sizes[i],196                )197            )198        self.patch_embeddings = nn.ModuleList(embeddings)199 200        # Transformer blocks201        blocks = []202        cur = 0203        for i in range(config.num_encoder_blocks):204            # each block consists of layers205            layers = []206            if i != 0:207                cur += config.depths[i - 1]208            for j in range(config.depths[i]):209                layers.append(210                    PoolFormerLayer(211                        config,212                        num_channels=config.hidden_sizes[i],213                        pool_size=config.pool_size,214                        hidden_size=config.hidden_sizes[i],215                        intermediate_size=int(config.hidden_sizes[i] * config.mlp_ratio),216                        drop_path=dpr[cur + j],217                    )218                )219            blocks.append(nn.ModuleList(layers))220 221        self.block = nn.ModuleList(blocks)222 223    def forward(self, pixel_values, output_hidden_states=False, return_dict=True):224        all_hidden_states = () if output_hidden_states else None225 226        hidden_states = pixel_values227        for idx, layers in enumerate(zip(self.patch_embeddings, self.block)):228            embedding_layer, block_layer = layers229            # Get patch embeddings from hidden_states230            hidden_states = embedding_layer(hidden_states)231            # Send the embeddings through the blocks232            for _, blk in enumerate(block_layer):233                layer_outputs = blk(hidden_states)234                hidden_states = layer_outputs[0]235 236            if output_hidden_states:237                all_hidden_states = all_hidden_states + (hidden_states,)238 239        if not return_dict:240            return tuple(v for v in [hidden_states, all_hidden_states] if v is not None)241 242        return BaseModelOutputWithNoAttention(last_hidden_state=hidden_states, hidden_states=all_hidden_states)243 244 245@auto_docstring246class PoolFormerPreTrainedModel(PreTrainedModel):247    config: PoolFormerConfig248    base_model_prefix = "poolformer"249    main_input_name = "pixel_values"250    _no_split_modules = ["PoolFormerLayer"]251 252    def _init_weights(self, module):253        """Initialize the weights"""254        if isinstance(module, (nn.Linear, nn.Conv2d)):255            module.weight.data.normal_(mean=0.0, std=self.config.initializer_range)256            if module.bias is not None:257                module.bias.data.zero_()258        elif isinstance(module, nn.GroupNorm):259            module.bias.data.zero_()260            module.weight.data.fill_(1.0)261        elif isinstance(module, PoolFormerLayer):262            if hasattr(module, "layer_scale_1"):263                module.layer_scale_1.data.fill_(self.config.layer_scale_init_value)264                module.layer_scale_2.data.fill_(self.config.layer_scale_init_value)265 266 267@auto_docstring268class PoolFormerModel(PoolFormerPreTrainedModel):269    def __init__(self, config):270        super().__init__(config)271        self.config = config272 273        self.encoder = PoolFormerEncoder(config)274 275        # Initialize weights and apply final processing276        self.post_init()277 278    def get_input_embeddings(self):279        return self.embeddings.patch_embeddings280 281    @auto_docstring282    def forward(283        self,284        pixel_values: Optional[torch.FloatTensor] = None,285        output_hidden_states: Optional[bool] = None,286        return_dict: Optional[bool] = None,287    ) -> Union[tuple, BaseModelOutputWithNoAttention]:288        output_hidden_states = (289            output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states290        )291        return_dict = return_dict if return_dict is not None else self.config.use_return_dict292 293        if pixel_values is None:294            raise ValueError("You have to specify pixel_values")295 296        encoder_outputs = self.encoder(297            pixel_values,298            output_hidden_states=output_hidden_states,299            return_dict=return_dict,300        )301        sequence_output = encoder_outputs[0]302 303        if not return_dict:304            return (sequence_output, None) + encoder_outputs[1:]305 306        return BaseModelOutputWithNoAttention(307            last_hidden_state=sequence_output,308            hidden_states=encoder_outputs.hidden_states,309        )310 311 312class PoolFormerFinalPooler(nn.Module):313    def __init__(self, config):314        super().__init__()315        self.dense = nn.Linear(config.hidden_size, config.hidden_size)316 317    def forward(self, hidden_states):318        output = self.dense(hidden_states)319        return output320 321 322@auto_docstring(323    custom_intro="""324    PoolFormer Model transformer with an image classification head on top325    """326)327class PoolFormerForImageClassification(PoolFormerPreTrainedModel):328    def __init__(self, config):329        super().__init__(config)330        self.num_labels = config.num_labels331        self.poolformer = PoolFormerModel(config)332 333        # Final norm334        self.norm = PoolFormerGroupNorm(config.hidden_sizes[-1])335        # Classifier head336        self.classifier = (337            nn.Linear(config.hidden_sizes[-1], config.num_labels) if config.num_labels > 0 else nn.Identity()338        )339 340        # Initialize weights and apply final processing341        self.post_init()342 343    @auto_docstring344    def forward(345        self,346        pixel_values: Optional[torch.FloatTensor] = None,347        labels: Optional[torch.LongTensor] = None,348        output_hidden_states: Optional[bool] = None,349        return_dict: Optional[bool] = None,350    ) -> Union[tuple, ImageClassifierOutputWithNoAttention]:351        r"""352        labels (`torch.LongTensor` of shape `(batch_size,)`, *optional*):353            Labels for computing the image classification/regression loss. Indices should be in `[0, ...,354            config.num_labels - 1]`. If `config.num_labels == 1` a regression loss is computed (Mean-Square loss), If355            `config.num_labels > 1` a classification loss is computed (Cross-Entropy).356        """357        return_dict = return_dict if return_dict is not None else self.config.use_return_dict358 359        outputs = self.poolformer(360            pixel_values,361            output_hidden_states=output_hidden_states,362            return_dict=return_dict,363        )364 365        sequence_output = outputs[0]366 367        logits = self.classifier(self.norm(sequence_output).mean([-2, -1]))368 369        loss = None370        if labels is not None:371            loss = self.loss_function(labels, logits, self.config)372 373        if not return_dict:374            output = (logits,) + outputs[2:]375            return ((loss,) + output) if loss is not None else output376 377        return ImageClassifierOutputWithNoAttention(loss=loss, logits=logits, hidden_states=outputs.hidden_states)378 379 380__all__ = ["PoolFormerForImageClassification", "PoolFormerModel", "PoolFormerPreTrainedModel"]381