Aluode/PerceptionLabPortable

1# coding=utf-82# Copyright 2023 The Google Flax Team Authors and The HuggingFace Inc. team.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 16 17from functools import partial18from typing import Optional19 20import flax.linen as nn21import jax22import jax.numpy as jnp23from flax.core.frozen_dict import FrozenDict, freeze, unfreeze24from flax.traverse_util import flatten_dict, unflatten_dict25 26from transformers import RegNetConfig27from transformers.modeling_flax_outputs import (28    FlaxBaseModelOutputWithNoAttention,29    FlaxBaseModelOutputWithPooling,30    FlaxBaseModelOutputWithPoolingAndNoAttention,31    FlaxImageClassifierOutputWithNoAttention,32)33from transformers.modeling_flax_utils import (34    ACT2FN,35    FlaxPreTrainedModel,36    append_replace_return_docstrings,37    overwrite_call_docstring,38)39from transformers.utils import (40    add_start_docstrings,41    add_start_docstrings_to_model_forward,42)43 44 45REGNET_START_DOCSTRING = r"""46 47    This model inherits from [`FlaxPreTrainedModel`]. Check the superclass documentation for the generic methods the48    library implements for all its model (such as downloading, saving and converting weights from PyTorch models)49 50    This model is also a51    [flax.linen.Module](https://flax.readthedocs.io/en/latest/api_reference/flax.linen/module.html) subclass. Use it as52    a regular Flax linen Module and refer to the Flax documentation for all matter related to general usage and53    behavior.54 55    Finally, this model supports inherent JAX features such as:56 57    - [Just-In-Time (JIT) compilation](https://jax.readthedocs.io/en/latest/jax.html#just-in-time-compilation-jit)58    - [Automatic Differentiation](https://jax.readthedocs.io/en/latest/jax.html#automatic-differentiation)59    - [Vectorization](https://jax.readthedocs.io/en/latest/jax.html#vectorization-vmap)60    - [Parallelization](https://jax.readthedocs.io/en/latest/jax.html#parallelization-pmap)61 62    Parameters:63        config ([`RegNetConfig`]): Model configuration class with all the parameters of the model.64            Initializing with a config file does not load the weights associated with the model, only the65            configuration. Check out the [`~FlaxPreTrainedModel.from_pretrained`] method to load the model weights.66        dtype (`jax.numpy.dtype`, *optional*, defaults to `jax.numpy.float32`):67            The data type of the computation. Can be one of `jax.numpy.float32`, `jax.numpy.float16` (on GPUs) and68            `jax.numpy.bfloat16` (on TPUs).69 70            This can be used to enable mixed-precision training or half-precision inference on GPUs or TPUs. If71            specified all the computation will be performed with the given `dtype`.72 73            **Note that this only specifies the dtype of the computation and does not influence the dtype of model74            parameters.**75 76            If you wish to change the dtype of the model parameters, see [`~FlaxPreTrainedModel.to_fp16`] and77            [`~FlaxPreTrainedModel.to_bf16`].78"""79 80REGNET_INPUTS_DOCSTRING = r"""81    Args:82        pixel_values (`numpy.ndarray` of shape `(batch_size, num_channels, height, width)`):83            Pixel values. Pixel values can be obtained using [`AutoImageProcessor`]. See84            [`RegNetImageProcessor.__call__`] for details.85 86        output_hidden_states (`bool`, *optional*):87            Whether or not to return the hidden states of all layers. See `hidden_states` under returned tensors for88            more detail.89        return_dict (`bool`, *optional*):90            Whether or not to return a [`~utils.ModelOutput`] instead of a plain tuple.91"""92 93 94# Copied from transformers.models.resnet.modeling_flax_resnet.Identity95class Identity(nn.Module):96    """Identity function."""97 98    @nn.compact99    def __call__(self, x, **kwargs):100        return x101 102 103class FlaxRegNetConvLayer(nn.Module):104    out_channels: int105    kernel_size: int = 3106    stride: int = 1107    groups: int = 1108    activation: Optional[str] = "relu"109    dtype: jnp.dtype = jnp.float32110 111    def setup(self):112        self.convolution = nn.Conv(113            self.out_channels,114            kernel_size=(self.kernel_size, self.kernel_size),115            strides=self.stride,116            padding=self.kernel_size // 2,117            feature_group_count=self.groups,118            use_bias=False,119            kernel_init=nn.initializers.variance_scaling(2.0, mode="fan_out", distribution="truncated_normal"),120            dtype=self.dtype,121        )122        self.normalization = nn.BatchNorm(momentum=0.9, epsilon=1e-05, dtype=self.dtype)123        self.activation_func = ACT2FN[self.activation] if self.activation is not None else Identity()124 125    def __call__(self, hidden_state: jnp.ndarray, deterministic: bool = True) -> jnp.ndarray:126        hidden_state = self.convolution(hidden_state)127        hidden_state = self.normalization(hidden_state, use_running_average=deterministic)128        hidden_state = self.activation_func(hidden_state)129        return hidden_state130 131 132class FlaxRegNetEmbeddings(nn.Module):133    config: RegNetConfig134    dtype: jnp.dtype = jnp.float32135 136    def setup(self):137        self.embedder = FlaxRegNetConvLayer(138            self.config.embedding_size,139            kernel_size=3,140            stride=2,141            activation=self.config.hidden_act,142            dtype=self.dtype,143        )144 145    def __call__(self, pixel_values: jnp.ndarray, deterministic: bool = True) -> jnp.ndarray:146        num_channels = pixel_values.shape[-1]147        if num_channels != self.config.num_channels:148            raise ValueError(149                "Make sure that the channel dimension of the pixel values match with the one set in the configuration."150            )151        hidden_state = self.embedder(pixel_values, deterministic=deterministic)152        return hidden_state153 154 155# Copied from transformers.models.resnet.modeling_flax_resnet.FlaxResNetShortCut with ResNet->RegNet156class FlaxRegNetShortCut(nn.Module):157    """158    RegNet shortcut, used to project the residual features to the correct size. If needed, it is also used to159    downsample the input using `stride=2`.160    """161 162    out_channels: int163    stride: int = 2164    dtype: jnp.dtype = jnp.float32165 166    def setup(self):167        self.convolution = nn.Conv(168            self.out_channels,169            kernel_size=(1, 1),170            strides=self.stride,171            use_bias=False,172            kernel_init=nn.initializers.variance_scaling(2.0, mode="fan_out", distribution="truncated_normal"),173            dtype=self.dtype,174        )175        self.normalization = nn.BatchNorm(momentum=0.9, epsilon=1e-05, dtype=self.dtype)176 177    def __call__(self, x: jnp.ndarray, deterministic: bool = True) -> jnp.ndarray:178        hidden_state = self.convolution(x)179        hidden_state = self.normalization(hidden_state, use_running_average=deterministic)180        return hidden_state181 182 183class FlaxRegNetSELayerCollection(nn.Module):184    in_channels: int185    reduced_channels: int186    dtype: jnp.dtype = jnp.float32187 188    def setup(self):189        self.conv_1 = nn.Conv(190            self.reduced_channels,191            kernel_size=(1, 1),192            kernel_init=nn.initializers.variance_scaling(2.0, mode="fan_out", distribution="truncated_normal"),193            dtype=self.dtype,194            name="0",195        )  # 0 is the name used in corresponding pytorch implementation196        self.conv_2 = nn.Conv(197            self.in_channels,198            kernel_size=(1, 1),199            kernel_init=nn.initializers.variance_scaling(2.0, mode="fan_out", distribution="truncated_normal"),200            dtype=self.dtype,201            name="2",202        )  # 2 is the name used in corresponding pytorch implementation203 204    def __call__(self, hidden_state: jnp.ndarray) -> jnp.ndarray:205        hidden_state = self.conv_1(hidden_state)206        hidden_state = nn.relu(hidden_state)207        hidden_state = self.conv_2(hidden_state)208        attention = nn.sigmoid(hidden_state)209 210        return attention211 212 213class FlaxRegNetSELayer(nn.Module):214    """215    Squeeze and Excitation layer (SE) proposed in [Squeeze-and-Excitation Networks](https://huggingface.co/papers/1709.01507).216    """217 218    in_channels: int219    reduced_channels: int220    dtype: jnp.dtype = jnp.float32221 222    def setup(self):223        self.pooler = partial(nn.avg_pool, padding=((0, 0), (0, 0)))224        self.attention = FlaxRegNetSELayerCollection(self.in_channels, self.reduced_channels, dtype=self.dtype)225 226    def __call__(self, hidden_state: jnp.ndarray) -> jnp.ndarray:227        pooled = self.pooler(228            hidden_state,229            window_shape=(hidden_state.shape[1], hidden_state.shape[2]),230            strides=(hidden_state.shape[1], hidden_state.shape[2]),231        )232        attention = self.attention(pooled)233        hidden_state = hidden_state * attention234        return hidden_state235 236 237class FlaxRegNetXLayerCollection(nn.Module):238    config: RegNetConfig239    out_channels: int240    stride: int = 1241    dtype: jnp.dtype = jnp.float32242 243    def setup(self):244        groups = max(1, self.out_channels // self.config.groups_width)245 246        self.layer = [247            FlaxRegNetConvLayer(248                self.out_channels,249                kernel_size=1,250                activation=self.config.hidden_act,251                dtype=self.dtype,252                name="0",253            ),254            FlaxRegNetConvLayer(255                self.out_channels,256                stride=self.stride,257                groups=groups,258                activation=self.config.hidden_act,259                dtype=self.dtype,260                name="1",261            ),262            FlaxRegNetConvLayer(263                self.out_channels,264                kernel_size=1,265                activation=None,266                dtype=self.dtype,267                name="2",268            ),269        ]270 271    def __call__(self, hidden_state: jnp.ndarray, deterministic: bool = True) -> jnp.ndarray:272        for layer in self.layer:273            hidden_state = layer(hidden_state, deterministic=deterministic)274        return hidden_state275 276 277class FlaxRegNetXLayer(nn.Module):278    """279    RegNet's layer composed by three `3x3` convolutions, same as a ResNet bottleneck layer with reduction = 1.280    """281 282    config: RegNetConfig283    in_channels: int284    out_channels: int285    stride: int = 1286    dtype: jnp.dtype = jnp.float32287 288    def setup(self):289        should_apply_shortcut = self.in_channels != self.out_channels or self.stride != 1290        self.shortcut = (291            FlaxRegNetShortCut(292                self.out_channels,293                stride=self.stride,294                dtype=self.dtype,295            )296            if should_apply_shortcut297            else Identity()298        )299        self.layer = FlaxRegNetXLayerCollection(300            self.config,301            in_channels=self.in_channels,302            out_channels=self.out_channels,303            stride=self.stride,304            dtype=self.dtype,305        )306        self.activation_func = ACT2FN[self.config.hidden_act]307 308    def __call__(self, hidden_state: jnp.ndarray, deterministic: bool = True) -> jnp.ndarray:309        residual = hidden_state310        hidden_state = self.layer(hidden_state)311        residual = self.shortcut(residual, deterministic=deterministic)312        hidden_state += residual313        hidden_state = self.activation_func(hidden_state)314        return hidden_state315 316 317class FlaxRegNetYLayerCollection(nn.Module):318    config: RegNetConfig319    in_channels: int320    out_channels: int321    stride: int = 1322    dtype: jnp.dtype = jnp.float32323 324    def setup(self):325        groups = max(1, self.out_channels // self.config.groups_width)326 327        self.layer = [328            FlaxRegNetConvLayer(329                self.out_channels,330                kernel_size=1,331                activation=self.config.hidden_act,332                dtype=self.dtype,333                name="0",334            ),335            FlaxRegNetConvLayer(336                self.out_channels,337                stride=self.stride,338                groups=groups,339                activation=self.config.hidden_act,340                dtype=self.dtype,341                name="1",342            ),343            FlaxRegNetSELayer(344                self.out_channels,345                reduced_channels=int(round(self.in_channels / 4)),346                dtype=self.dtype,347                name="2",348            ),349            FlaxRegNetConvLayer(350                self.out_channels,351                kernel_size=1,352                activation=None,353                dtype=self.dtype,354                name="3",355            ),356        ]357 358    def __call__(self, hidden_state: jnp.ndarray) -> jnp.ndarray:359        for layer in self.layer:360            hidden_state = layer(hidden_state)361        return hidden_state362 363 364class FlaxRegNetYLayer(nn.Module):365    """366    RegNet's Y layer: an X layer with Squeeze and Excitation.367    """368 369    config: RegNetConfig370    in_channels: int371    out_channels: int372    stride: int = 1373    dtype: jnp.dtype = jnp.float32374 375    def setup(self):376        should_apply_shortcut = self.in_channels != self.out_channels or self.stride != 1377 378        self.shortcut = (379            FlaxRegNetShortCut(380                self.out_channels,381                stride=self.stride,382                dtype=self.dtype,383            )384            if should_apply_shortcut385            else Identity()386        )387        self.layer = FlaxRegNetYLayerCollection(388            self.config,389            in_channels=self.in_channels,390            out_channels=self.out_channels,391            stride=self.stride,392            dtype=self.dtype,393        )394        self.activation_func = ACT2FN[self.config.hidden_act]395 396    def __call__(self, hidden_state: jnp.ndarray, deterministic: bool = True) -> jnp.ndarray:397        residual = hidden_state398        hidden_state = self.layer(hidden_state)399        residual = self.shortcut(residual, deterministic=deterministic)400        hidden_state += residual401        hidden_state = self.activation_func(hidden_state)402        return hidden_state403 404 405class FlaxRegNetStageLayersCollection(nn.Module):406    """407    A RegNet stage composed by stacked layers.408    """409 410    config: RegNetConfig411    in_channels: int412    out_channels: int413    stride: int = 2414    depth: int = 2415    dtype: jnp.dtype = jnp.float32416 417    def setup(self):418        layer = FlaxRegNetXLayer if self.config.layer_type == "x" else FlaxRegNetYLayer419 420        layers = [421            # downsampling is done in the first layer with stride of 2422            layer(423                self.config,424                self.in_channels,425                self.out_channels,426                stride=self.stride,427                dtype=self.dtype,428                name="0",429            )430        ]431 432        for i in range(self.depth - 1):433            layers.append(434                layer(435                    self.config,436                    self.out_channels,437                    self.out_channels,438                    dtype=self.dtype,439                    name=str(i + 1),440                )441            )442 443        self.layers = layers444 445    def __call__(self, x: jnp.ndarray, deterministic: bool = True) -> jnp.ndarray:446        hidden_state = x447        for layer in self.layers:448            hidden_state = layer(hidden_state, deterministic=deterministic)449        return hidden_state450 451 452# Copied from transformers.models.resnet.modeling_flax_resnet.FlaxResNetStage with ResNet->RegNet453class FlaxRegNetStage(nn.Module):454    """455    A RegNet stage composed by stacked layers.456    """457 458    config: RegNetConfig459    in_channels: int460    out_channels: int461    stride: int = 2462    depth: int = 2463    dtype: jnp.dtype = jnp.float32464 465    def setup(self):466        self.layers = FlaxRegNetStageLayersCollection(467            self.config,468            in_channels=self.in_channels,469            out_channels=self.out_channels,470            stride=self.stride,471            depth=self.depth,472            dtype=self.dtype,473        )474 475    def __call__(self, x: jnp.ndarray, deterministic: bool = True) -> jnp.ndarray:476        return self.layers(x, deterministic=deterministic)477 478 479# Copied from transformers.models.resnet.modeling_flax_resnet.FlaxResNetStageCollection with ResNet->RegNet480class FlaxRegNetStageCollection(nn.Module):481    config: RegNetConfig482    dtype: jnp.dtype = jnp.float32483 484    def setup(self):485        in_out_channels = zip(self.config.hidden_sizes, self.config.hidden_sizes[1:])486        stages = [487            FlaxRegNetStage(488                self.config,489                self.config.embedding_size,490                self.config.hidden_sizes[0],491                stride=2 if self.config.downsample_in_first_stage else 1,492                depth=self.config.depths[0],493                dtype=self.dtype,494                name="0",495            )496        ]497 498        for i, ((in_channels, out_channels), depth) in enumerate(zip(in_out_channels, self.config.depths[1:])):499            stages.append(500                FlaxRegNetStage(self.config, in_channels, out_channels, depth=depth, dtype=self.dtype, name=str(i + 1))501            )502 503        self.stages = stages504 505    def __call__(506        self,507        hidden_state: jnp.ndarray,508        output_hidden_states: bool = False,509        deterministic: bool = True,510    ) -> FlaxBaseModelOutputWithNoAttention:511        hidden_states = () if output_hidden_states else None512 513        for stage_module in self.stages:514            if output_hidden_states:515                hidden_states = hidden_states + (hidden_state.transpose(0, 3, 1, 2),)516 517            hidden_state = stage_module(hidden_state, deterministic=deterministic)518 519        return hidden_state, hidden_states520 521 522# Copied from transformers.models.resnet.modeling_flax_resnet.FlaxResNetEncoder with ResNet->RegNet523class FlaxRegNetEncoder(nn.Module):524    config: RegNetConfig525    dtype: jnp.dtype = jnp.float32526 527    def setup(self):528        self.stages = FlaxRegNetStageCollection(self.config, dtype=self.dtype)529 530    def __call__(531        self,532        hidden_state: jnp.ndarray,533        output_hidden_states: bool = False,534        return_dict: bool = True,535        deterministic: bool = True,536    ) -> FlaxBaseModelOutputWithNoAttention:537        hidden_state, hidden_states = self.stages(538            hidden_state, output_hidden_states=output_hidden_states, deterministic=deterministic539        )540 541        if output_hidden_states:542            hidden_states = hidden_states + (hidden_state.transpose(0, 3, 1, 2),)543 544        if not return_dict:545            return tuple(v for v in [hidden_state, hidden_states] if v is not None)546 547        return FlaxBaseModelOutputWithNoAttention(548            last_hidden_state=hidden_state,549            hidden_states=hidden_states,550        )551 552 553# Copied from transformers.models.resnet.modeling_flax_resnet.FlaxResNetPreTrainedModel with ResNet->RegNet,resnet->regnet,RESNET->REGNET554class FlaxRegNetPreTrainedModel(FlaxPreTrainedModel):555    """556    An abstract class to handle weights initialization and a simple interface for downloading and loading pretrained557    models.558    """559 560    config_class = RegNetConfig561    base_model_prefix = "regnet"562    main_input_name = "pixel_values"563    module_class: nn.Module = None564 565    def __init__(566        self,567        config: RegNetConfig,568        input_shape=(1, 224, 224, 3),569        seed: int = 0,570        dtype: jnp.dtype = jnp.float32,571        _do_init: bool = True,572        **kwargs,573    ):574        module = self.module_class(config=config, dtype=dtype, **kwargs)575        if input_shape is None:576            input_shape = (1, config.image_size, config.image_size, config.num_channels)577        super().__init__(config, module, input_shape=input_shape, seed=seed, dtype=dtype, _do_init=_do_init)578 579    def init_weights(self, rng: jax.random.PRNGKey, input_shape: tuple, params: FrozenDict = None) -> FrozenDict:580        # init input tensors581        pixel_values = jnp.zeros(input_shape, dtype=self.dtype)582 583        rngs = {"params": rng}584 585        random_params = self.module.init(rngs, pixel_values, return_dict=False)586 587        if params is not None:588            random_params = flatten_dict(unfreeze(random_params))589            params = flatten_dict(unfreeze(params))590            for missing_key in self._missing_keys:591                params[missing_key] = random_params[missing_key]592            self._missing_keys = set()593            return freeze(unflatten_dict(params))594        else:595            return random_params596 597    @add_start_docstrings_to_model_forward(REGNET_INPUTS_DOCSTRING)598    def __call__(599        self,600        pixel_values,601        params: Optional[dict] = None,602        train: bool = False,603        output_hidden_states: Optional[bool] = None,604        return_dict: Optional[bool] = None,605    ):606        output_hidden_states = (607            output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states608        )609        return_dict = return_dict if return_dict is not None else self.config.return_dict610 611        pixel_values = jnp.transpose(pixel_values, (0, 2, 3, 1))612 613        # Handle any PRNG if needed614        rngs = {}615 616        return self.module.apply(617            {618                "params": params["params"] if params is not None else self.params["params"],619                "batch_stats": params["batch_stats"] if params is not None else self.params["batch_stats"],620            },621            jnp.array(pixel_values, dtype=jnp.float32),622            not train,623            output_hidden_states,624            return_dict,625            rngs=rngs,626            mutable=["batch_stats"] if train else False,  # Returning tuple with batch_stats only when train is True627        )628 629 630# Copied from transformers.models.resnet.modeling_flax_resnet.FlaxResNetModule with ResNet->RegNet631class FlaxRegNetModule(nn.Module):632    config: RegNetConfig633    dtype: jnp.dtype = jnp.float32  # the dtype of the computation634 635    def setup(self):636        self.embedder = FlaxRegNetEmbeddings(self.config, dtype=self.dtype)637        self.encoder = FlaxRegNetEncoder(self.config, dtype=self.dtype)638 639        # Adaptive average pooling used in resnet640        self.pooler = partial(641            nn.avg_pool,642            padding=((0, 0), (0, 0)),643        )644 645    def __call__(646        self,647        pixel_values,648        deterministic: bool = True,649        output_hidden_states: bool = False,650        return_dict: bool = True,651    ) -> FlaxBaseModelOutputWithPoolingAndNoAttention:652        output_hidden_states = (653            output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states654        )655        return_dict = return_dict if return_dict is not None else self.config.use_return_dict656 657        embedding_output = self.embedder(pixel_values, deterministic=deterministic)658 659        encoder_outputs = self.encoder(660            embedding_output,661            output_hidden_states=output_hidden_states,662            return_dict=return_dict,663            deterministic=deterministic,664        )665 666        last_hidden_state = encoder_outputs[0]667 668        pooled_output = self.pooler(669            last_hidden_state,670            window_shape=(last_hidden_state.shape[1], last_hidden_state.shape[2]),671            strides=(last_hidden_state.shape[1], last_hidden_state.shape[2]),672        ).transpose(0, 3, 1, 2)673 674        last_hidden_state = last_hidden_state.transpose(0, 3, 1, 2)675 676        if not return_dict:677            return (last_hidden_state, pooled_output) + encoder_outputs[1:]678 679        return FlaxBaseModelOutputWithPoolingAndNoAttention(680            last_hidden_state=last_hidden_state,681            pooler_output=pooled_output,682            hidden_states=encoder_outputs.hidden_states,683        )684 685 686@add_start_docstrings(687    "The bare RegNet model outputting raw features without any specific head on top.",688    REGNET_START_DOCSTRING,689)690class FlaxRegNetModel(FlaxRegNetPreTrainedModel):691    module_class = FlaxRegNetModule692 693 694FLAX_VISION_MODEL_DOCSTRING = """695    Returns:696 697    Examples:698 699    ```python700    >>> from transformers import AutoImageProcessor, FlaxRegNetModel701    >>> from PIL import Image702    >>> import requests703 704    >>> url = "http://images.cocodataset.org/val2017/000000039769.jpg"705    >>> image = Image.open(requests.get(url, stream=True).raw)706 707    >>> image_processor = AutoImageProcessor.from_pretrained("facebook/regnet-y-040")708    >>> model = FlaxRegNetModel.from_pretrained("facebook/regnet-y-040")709 710    >>> inputs = image_processor(images=image, return_tensors="np")711    >>> outputs = model(**inputs)712    >>> last_hidden_states = outputs.last_hidden_state713    ```714"""715 716overwrite_call_docstring(FlaxRegNetModel, FLAX_VISION_MODEL_DOCSTRING)717append_replace_return_docstrings(718    FlaxRegNetModel,719    output_type=FlaxBaseModelOutputWithPooling,720    config_class=RegNetConfig,721)722 723 724# Copied from transformers.models.resnet.modeling_flax_resnet.FlaxResNetClassifierCollection with ResNet->RegNet725class FlaxRegNetClassifierCollection(nn.Module):726    config: RegNetConfig727    dtype: jnp.dtype = jnp.float32728 729    def setup(self):730        self.classifier = nn.Dense(self.config.num_labels, dtype=self.dtype, name="1")731 732    def __call__(self, x: jnp.ndarray) -> jnp.ndarray:733        return self.classifier(x)734 735 736# Copied from transformers.models.resnet.modeling_flax_resnet.FlaxResNetForImageClassificationModule with ResNet->RegNet,resnet->regnet,RESNET->REGNET737class FlaxRegNetForImageClassificationModule(nn.Module):738    config: RegNetConfig739    dtype: jnp.dtype = jnp.float32740 741    def setup(self):742        self.regnet = FlaxRegNetModule(config=self.config, dtype=self.dtype)743 744        if self.config.num_labels > 0:745            self.classifier = FlaxRegNetClassifierCollection(self.config, dtype=self.dtype)746        else:747            self.classifier = Identity()748 749    def __call__(750        self,751        pixel_values=None,752        deterministic: bool = True,753        output_hidden_states=None,754        return_dict=None,755    ):756        return_dict = return_dict if return_dict is not None else self.config.use_return_dict757 758        outputs = self.regnet(759            pixel_values,760            deterministic=deterministic,761            output_hidden_states=output_hidden_states,762            return_dict=return_dict,763        )764 765        pooled_output = outputs.pooler_output if return_dict else outputs[1]766 767        logits = self.classifier(pooled_output[:, :, 0, 0])768 769        if not return_dict:770            output = (logits,) + outputs[2:]771            return output772 773        return FlaxImageClassifierOutputWithNoAttention(logits=logits, hidden_states=outputs.hidden_states)774 775 776@add_start_docstrings(777    """778    RegNet Model with an image classification head on top (a linear layer on top of the pooled features), e.g. for779    ImageNet.780    """,781    REGNET_START_DOCSTRING,782)783class FlaxRegNetForImageClassification(FlaxRegNetPreTrainedModel):784    module_class = FlaxRegNetForImageClassificationModule785 786 787FLAX_VISION_CLASSIF_DOCSTRING = """788    Returns:789 790    Example:791 792    ```python793    >>> from transformers import AutoImageProcessor, FlaxRegNetForImageClassification794    >>> from PIL import Image795    >>> import jax796    >>> import requests797 798    >>> url = "http://images.cocodataset.org/val2017/000000039769.jpg"799    >>> image = Image.open(requests.get(url, stream=True).raw)800 801    >>> image_processor = AutoImageProcessor.from_pretrained("facebook/regnet-y-040")802    >>> model = FlaxRegNetForImageClassification.from_pretrained("facebook/regnet-y-040")803 804    >>> inputs = image_processor(images=image, return_tensors="np")805    >>> outputs = model(**inputs)806    >>> logits = outputs.logits807 808    >>> # model predicts one of the 1000 ImageNet classes809    >>> predicted_class_idx = jax.numpy.argmax(logits, axis=-1)810    >>> print("Predicted class:", model.config.id2label[predicted_class_idx.item()])811    ```812"""813 814overwrite_call_docstring(FlaxRegNetForImageClassification, FLAX_VISION_CLASSIF_DOCSTRING)815append_replace_return_docstrings(816    FlaxRegNetForImageClassification,817    output_type=FlaxImageClassifierOutputWithNoAttention,818    config_class=RegNetConfig,819)820 821 822__all__ = ["FlaxRegNetForImageClassification", "FlaxRegNetModel", "FlaxRegNetPreTrainedModel"]823