Aluode/PerceptionLabPortable
0
1# coding=utf-82# Copyright 2022 Google AI 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 BiT model. Also supports backbone for ViT hybrid."""16 17import collections18import math19from typing import Optional20 21import numpy as np22import torch23from torch import Tensor, nn24 25from ...activations import ACT2FN26from ...modeling_outputs import (27 BackboneOutput,28 BaseModelOutputWithNoAttention,29 BaseModelOutputWithPoolingAndNoAttention,30 ImageClassifierOutputWithNoAttention,31)32from ...modeling_utils import PreTrainedModel33from ...utils import auto_docstring, logging34from ...utils.backbone_utils import BackboneMixin35from .configuration_bit import BitConfig36 37 38logger = logging.get_logger(__name__)39 40 41def get_padding_value(padding=None, kernel_size=7, stride=1, dilation=1) -> tuple[tuple, bool]:42 r"""43 Utility function to get the tuple padding value given the kernel_size and padding.44 45 Args:46 padding (Union[`str`, `int`], *optional*):47 Padding value, can be either `"same"`, `"valid"`. If a different value is provided the default padding from48 PyTorch is used.49 kernel_size (`int`, *optional*, defaults to 7):50 Kernel size of the convolution layers.51 stride (`int`, *optional*, defaults to 1):52 Stride value of the convolution layers.53 dilation (`int`, *optional*, defaults to 1):54 Dilation value of the convolution layers.55 """56 dynamic = False57 if padding is None:58 padding = ((stride - 1) + dilation * (kernel_size - 1)) // 259 return padding, dynamic60 61 if isinstance(padding, str):62 # for any string padding, the padding will be calculated for you, one of three ways63 padding = padding.lower()64 if padding == "same":65 # TF compatible 'SAME' padding, has a performance and GPU memory allocation impact66 if stride == 1 and (dilation * (kernel_size - 1)) % 2 == 0:67 # static case, no extra overhead68 padding = ((stride - 1) + dilation * (kernel_size - 1)) // 269 else:70 # dynamic 'SAME' padding, has runtime/GPU memory overhead71 padding = 072 dynamic = True73 elif padding == "valid":74 # 'VALID' padding, same as padding=075 padding = 076 else:77 # Default to PyTorch style 'same'-ish symmetric padding78 padding = ((stride - 1) + dilation * (kernel_size - 1)) // 279 return padding, dynamic80 81 82class WeightStandardizedConv2d(nn.Conv2d):83 """Conv2d with Weight Standardization. Includes TensorFlow compatible SAME padding. Used for ViT Hybrid model.84 85 Paper: [Micro-Batch Training with Batch-Channel Normalization and Weight86 Standardization](https://huggingface.co/papers/1903.10520v2)87 """88 89 def __init__(90 self,91 in_channel,92 out_channels,93 kernel_size,94 stride=1,95 padding="SAME",96 dilation=1,97 groups=1,98 bias=False,99 eps=1e-6,100 ):101 padding, is_dynamic = get_padding_value(padding, kernel_size, stride=stride, dilation=dilation)102 super().__init__(103 in_channel,104 out_channels,105 kernel_size,106 stride=stride,107 padding=padding,108 dilation=dilation,109 groups=groups,110 bias=bias,111 )112 if is_dynamic:113 self.pad = DynamicPad2d(kernel_size, stride, dilation)114 else:115 self.pad = None116 self.eps = eps117 118 def forward(self, hidden_state):119 if self.pad is not None:120 hidden_state = self.pad(hidden_state)121 weight = nn.functional.batch_norm(122 self.weight.reshape(1, self.out_channels, -1), None, None, training=True, momentum=0.0, eps=self.eps123 ).reshape_as(self.weight)124 hidden_state = nn.functional.conv2d(125 hidden_state, weight, self.bias, self.stride, self.padding, self.dilation, self.groups126 )127 return hidden_state128 129 130class BitGroupNormActivation(nn.GroupNorm):131 r"""132 A module that combines group normalization with an activation function.133 """134 135 def __init__(self, config, num_channels, eps=1e-5, affine=True, apply_activation=True):136 super().__init__(config.num_groups, num_channels, eps=eps, affine=affine)137 if apply_activation:138 self.activation = ACT2FN[config.hidden_act]139 else:140 self.activation = nn.Identity()141 142 def forward(self, hidden_state):143 hidden_state = nn.functional.group_norm(hidden_state, self.num_groups, self.weight, self.bias, self.eps)144 hidden_state = self.activation(hidden_state)145 return hidden_state146 147 148class DynamicPad2d(nn.Module):149 r"""150 A module that wraps dynamic padding of any input, given the parameters of the convolutional layer and the input151 hidden states.152 """153 154 def __init__(self, kernel_size, stride, dilation, value=0):155 super().__init__()156 # Safety checkers157 if isinstance(kernel_size, int):158 kernel_size = (kernel_size, kernel_size)159 160 if isinstance(stride, int):161 stride = (stride, stride)162 163 if isinstance(dilation, int):164 dilation = (dilation, dilation)165 166 self.kernel_size = kernel_size167 self.stride = stride168 self.dilation = dilation169 self.value = value170 171 def compute_padding(x, kernel_size, stride, dilation):172 return max((math.ceil(x / stride) - 1) * stride + (kernel_size - 1) * dilation + 1 - x, 0)173 174 self.compute_padding = compute_padding175 176 def forward(self, input):177 # Get width and height178 input_height, input_width = input.size()[-2:]179 180 # Compute the padding values181 padding_height = self.compute_padding(input_height, self.kernel_size[0], self.stride[0], self.dilation[0])182 padding_width = self.compute_padding(input_width, self.kernel_size[1], self.stride[1], self.dilation[1])183 184 # apply pad185 if padding_height > 0 or padding_width > 0:186 input = nn.functional.pad(187 input,188 [189 padding_width // 2,190 padding_width - padding_width // 2,191 padding_height // 2,192 padding_height - padding_height // 2,193 ],194 value=self.value,195 )196 return input197 198 199class BitMaxPool2d(nn.MaxPool2d):200 """Tensorflow like 'SAME' wrapper for 2D max pooling"""201 202 def __init__(203 self,204 kernel_size: int,205 stride=None,206 dilation=1,207 ceil_mode=False,208 padding=(0, 0),209 padding_value=0,210 use_dynamic_padding=True,211 ):212 kernel_size = kernel_size if isinstance(kernel_size, collections.abc.Iterable) else (kernel_size, kernel_size)213 stride = stride if isinstance(stride, collections.abc.Iterable) else (stride, stride)214 dilation = dilation if isinstance(dilation, collections.abc.Iterable) else (dilation, dilation)215 super().__init__(kernel_size, stride, padding, dilation, ceil_mode)216 if use_dynamic_padding:217 self.pad = DynamicPad2d(kernel_size, stride, dilation, padding_value)218 else:219 self.pad = nn.Identity()220 221 def forward(self, hidden_states):222 hidden_states = self.pad(hidden_states)223 return nn.functional.max_pool2d(224 hidden_states, self.kernel_size, self.stride, self.padding, self.dilation, self.ceil_mode225 )226 227 228class BitEmbeddings(nn.Module):229 """230 BiT Embeddings (stem) composed of a single aggressive convolution.231 """232 233 def __init__(self, config: BitConfig):234 super().__init__()235 236 self.convolution = WeightStandardizedConv2d(237 config.num_channels,238 config.embedding_size,239 kernel_size=7,240 stride=2,241 eps=1e-8,242 padding=config.global_padding,243 )244 245 self.pooler = BitMaxPool2d(kernel_size=3, stride=2, use_dynamic_padding=config.embedding_dynamic_padding)246 247 # Use the same padding strategy as convolutional layers248 if config.global_padding is not None and config.global_padding.upper() == "SAME":249 self.pad = nn.Identity()250 else:251 self.pad = nn.ConstantPad2d(padding=(1, 1, 1, 1), value=0.0)252 253 if config.layer_type != "preactivation":254 self.norm = BitGroupNormActivation(config, num_channels=config.embedding_size)255 else:256 self.norm = nn.Identity()257 258 self.num_channels = config.num_channels259 260 def forward(self, pixel_values: Tensor) -> Tensor:261 num_channels = pixel_values.shape[1]262 if num_channels != self.num_channels:263 raise ValueError(264 "Make sure that the channel dimension of the pixel values match with the one set in the configuration."265 )266 267 embedding = self.convolution(pixel_values)268 269 embedding = self.pad(embedding)270 271 embedding = self.norm(embedding)272 273 embedding = self.pooler(embedding)274 275 return embedding276 277 278# Copied from transformers.models.convnext.modeling_convnext.drop_path279def drop_path(input: torch.Tensor, drop_prob: float = 0.0, training: bool = False) -> torch.Tensor:280 """281 Drop paths (Stochastic Depth) per sample (when applied in main path of residual blocks).282 283 Comment by Ross Wightman: This is the same as the DropConnect impl I created for EfficientNet, etc networks,284 however, the original name is misleading as 'Drop Connect' is a different form of dropout in a separate paper...285 See discussion: https://github.com/tensorflow/tpu/issues/494#issuecomment-532968956 ... I've opted for changing the286 layer and argument names to 'drop path' rather than mix DropConnect as a layer name and use 'survival rate' as the287 argument.288 """289 if drop_prob == 0.0 or not training:290 return input291 keep_prob = 1 - drop_prob292 shape = (input.shape[0],) + (1,) * (input.ndim - 1) # work with diff dim tensors, not just 2D ConvNets293 random_tensor = keep_prob + torch.rand(shape, dtype=input.dtype, device=input.device)294 random_tensor.floor_() # binarize295 output = input.div(keep_prob) * random_tensor296 return output297 298 299# Copied from transformers.models.beit.modeling_beit.BeitDropPath with Beit->Bit300class BitDropPath(nn.Module):301 """Drop paths (Stochastic Depth) per sample (when applied in main path of residual blocks)."""302 303 def __init__(self, drop_prob: Optional[float] = None) -> None:304 super().__init__()305 self.drop_prob = drop_prob306 307 def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:308 return drop_path(hidden_states, self.drop_prob, self.training)309 310 def extra_repr(self) -> str:311 return f"p={self.drop_prob}"312 313 314def make_div(value, divisor=8):315 min_value = divisor316 new_value = max(min_value, int(value + divisor / 2) // divisor * divisor)317 if new_value < 0.9 * value:318 new_value += divisor319 return new_value320 321 322class BitPreActivationBottleneckLayer(nn.Module):323 """Pre-activation (v2) bottleneck block.324 Follows the implementation of "Identity Mappings in Deep Residual Networks":325 https://github.com/KaimingHe/resnet-1k-layers/blob/master/resnet-pre-act.lua326 327 Except it puts the stride on 3x3 conv when available.328 """329 330 def __init__(331 self,332 config,333 in_channels,334 out_channels=None,335 bottle_ratio=0.25,336 stride=1,337 dilation=1,338 first_dilation=None,339 groups=1,340 drop_path_rate=0.0,341 is_first_layer=False,342 ):343 super().__init__()344 345 first_dilation = first_dilation or dilation346 347 out_channels = out_channels or in_channels348 mid_channels = make_div(out_channels * bottle_ratio)349 350 if is_first_layer:351 self.downsample = BitDownsampleConv(352 config,353 in_channels,354 out_channels,355 stride=stride,356 preact=True,357 )358 else:359 self.downsample = None360 361 self.norm1 = BitGroupNormActivation(config, in_channels)362 self.conv1 = WeightStandardizedConv2d(in_channels, mid_channels, 1, eps=1e-8, padding=config.global_padding)363 364 self.norm2 = BitGroupNormActivation(config, num_channels=mid_channels)365 self.conv2 = WeightStandardizedConv2d(366 mid_channels, mid_channels, 3, stride=stride, groups=groups, eps=1e-8, padding=config.global_padding367 )368 369 self.norm3 = BitGroupNormActivation(config, mid_channels)370 self.conv3 = WeightStandardizedConv2d(mid_channels, out_channels, 1, eps=1e-8, padding=config.global_padding)371 372 self.drop_path = BitDropPath(drop_path_rate) if drop_path_rate > 0 else nn.Identity()373 374 def forward(self, hidden_states):375 hidden_states_preact = self.norm1(hidden_states)376 377 # shortcut branch378 shortcut = hidden_states379 if self.downsample is not None:380 shortcut = self.downsample(hidden_states_preact)381 382 # residual branch383 hidden_states = self.conv1(hidden_states_preact)384 hidden_states = self.conv2(self.norm2(hidden_states))385 hidden_states = self.conv3(self.norm3(hidden_states))386 hidden_states = self.drop_path(hidden_states)387 return hidden_states + shortcut388 389 390class BitBottleneckLayer(nn.Module):391 """Non Pre-activation bottleneck block, equivalent to V1.5/V1b bottleneck. Used for ViT Hybrid."""392 393 def __init__(394 self,395 config,396 in_channels,397 out_channels=None,398 bottle_ratio=0.25,399 stride=1,400 dilation=1,401 first_dilation=None,402 groups=1,403 drop_path_rate=0.0,404 is_first_layer=False,405 ):406 super().__init__()407 first_dilation = first_dilation or dilation408 409 out_channels = out_channels or in_channels410 mid_chs = make_div(out_channels * bottle_ratio)411 412 if is_first_layer:413 self.downsample = BitDownsampleConv(414 config,415 in_channels,416 out_channels,417 stride=stride,418 preact=False,419 )420 else:421 self.downsample = None422 423 self.conv1 = WeightStandardizedConv2d(in_channels, mid_chs, 1, eps=1e-8, padding=config.global_padding)424 self.norm1 = BitGroupNormActivation(config, num_channels=mid_chs)425 self.conv2 = WeightStandardizedConv2d(426 mid_chs,427 mid_chs,428 3,429 stride=stride,430 dilation=first_dilation,431 groups=groups,432 eps=1e-8,433 padding=config.global_padding,434 )435 self.norm2 = BitGroupNormActivation(config, num_channels=mid_chs)436 self.conv3 = WeightStandardizedConv2d(mid_chs, out_channels, 1, eps=1e-8, padding=config.global_padding)437 self.norm3 = BitGroupNormActivation(config, num_channels=out_channels, apply_activation=False)438 self.drop_path = BitDropPath(drop_path_rate) if drop_path_rate > 0 else nn.Identity()439 440 self.activation = ACT2FN[config.hidden_act]441 442 def forward(self, hidden_states):443 # shortcut branch444 shortcut = hidden_states445 if self.downsample is not None:446 shortcut = self.downsample(hidden_states)447 448 # residual449 hidden_states = self.conv1(hidden_states)450 hidden_states = self.norm1(hidden_states)451 452 hidden_states = self.conv2(hidden_states)453 hidden_states = self.norm2(hidden_states)454 455 hidden_states = self.conv3(hidden_states)456 hidden_states = self.norm3(hidden_states)457 458 hidden_states = self.drop_path(hidden_states)459 hidden_states = self.activation(hidden_states + shortcut)460 return hidden_states461 462 463class BitDownsampleConv(nn.Module):464 def __init__(465 self,466 config,467 in_channels,468 out_channels,469 stride=1,470 preact=True,471 ):472 super().__init__()473 self.conv = WeightStandardizedConv2d(474 in_channels, out_channels, 1, stride=stride, eps=1e-8, padding=config.global_padding475 )476 self.norm = (477 nn.Identity()478 if preact479 else BitGroupNormActivation(config, num_channels=out_channels, apply_activation=False)480 )481 482 def forward(self, x):483 return self.norm(self.conv(x))484 485 486class BitStage(nn.Module):487 """488 A ResNet v2 stage composed by stacked layers.489 """490 491 def __init__(492 self,493 config,494 in_channels,495 out_channels,496 stride,497 dilation,498 depth,499 bottle_ratio=0.25,500 layer_dropout=None,501 ):502 super().__init__()503 504 first_dilation = 1 if dilation in (1, 2) else 2505 506 # Get the layer type507 if config.layer_type == "bottleneck":508 layer_cls = BitBottleneckLayer509 else:510 layer_cls = BitPreActivationBottleneckLayer511 512 prev_chs = in_channels513 self.layers = nn.Sequential()514 for layer_idx in range(depth):515 # Get the current hyper-parameters516 stride, drop_path_rate, is_first_layer = self._get_updated_hyperparameters(517 layer_idx, stride, layer_dropout518 )519 520 self.layers.add_module(521 str(layer_idx),522 layer_cls(523 config,524 prev_chs,525 out_channels,526 stride=stride,527 dilation=dilation,528 bottle_ratio=bottle_ratio,529 first_dilation=first_dilation,530 drop_path_rate=drop_path_rate,531 is_first_layer=is_first_layer,532 ),533 )534 prev_chs = out_channels535 first_dilation = dilation536 537 def _get_updated_hyperparameters(self, layer_idx, stride, layer_dropout):538 r"""539 Get the new hyper-parameters with respect to the previous ones and the index of the current layer.540 """541 if layer_dropout:542 drop_path_rate = layer_dropout[layer_idx]543 else:544 drop_path_rate = 0.0545 546 if layer_idx != 0:547 stride = 1548 549 is_first_layer = layer_idx == 0550 551 return stride, drop_path_rate, is_first_layer552 553 def forward(self, input: Tensor) -> Tensor:554 hidden_state = input555 for _, layer in enumerate(self.layers):556 hidden_state = layer(hidden_state)557 return hidden_state558 559 560class BitEncoder(nn.Module):561 def __init__(self, config: BitConfig):562 super().__init__()563 self.stages = nn.ModuleList([])564 565 prev_chs = config.embedding_size566 567 # These needs to stay hardcoded568 current_stride = 4569 dilation = 1570 571 layer_dropouts = [572 x.tolist()573 for x in torch.Tensor(np.linspace(0, config.drop_path_rate, sum(config.depths))).split(config.depths)574 ]575 576 for stage_idx, (current_depth, current_hidden_size, layer_dropout) in enumerate(577 zip(config.depths, config.hidden_sizes, layer_dropouts)578 ):579 # Get the updated hyper params580 out_channels, stride, dilation = self._get_updated_hyperparameters(581 stage_idx, current_stride, current_hidden_size, dilation, config582 )583 584 stage = BitStage(585 config,586 prev_chs,587 out_channels,588 stride=stride,589 dilation=dilation,590 depth=current_depth,591 layer_dropout=layer_dropout,592 )593 594 prev_chs = out_channels595 current_stride *= stride596 597 self.stages.add_module(str(stage_idx), stage)598 599 def _get_updated_hyperparameters(self, stage_idx, current_stride, current_hidden_size, dilation, config):600 out_channels = make_div(current_hidden_size * config.width_factor)601 stride = 1 if stage_idx == 0 else 2602 if current_stride >= config.output_stride:603 dilation *= stride604 stride = 1605 return out_channels, stride, dilation606 607 def forward(608 self, hidden_state: Tensor, output_hidden_states: bool = False, return_dict: bool = True609 ) -> BaseModelOutputWithNoAttention:610 hidden_states = () if output_hidden_states else None611 612 for stage_module in self.stages:613 if output_hidden_states:614 hidden_states = hidden_states + (hidden_state,)615 616 hidden_state = stage_module(hidden_state)617 618 if output_hidden_states:619 hidden_states = hidden_states + (hidden_state,)620 621 if not return_dict:622 return tuple(v for v in [hidden_state, hidden_states] if v is not None)623 624 return BaseModelOutputWithNoAttention(625 last_hidden_state=hidden_state,626 hidden_states=hidden_states,627 )628 629 630@auto_docstring631class BitPreTrainedModel(PreTrainedModel):632 config: BitConfig633 base_model_prefix = "bit"634 main_input_name = "pixel_values"635 _no_split_modules = ["BitEmbeddings"]636 637 def _init_weights(self, module):638 if isinstance(module, nn.Conv2d):639 nn.init.kaiming_normal_(module.weight, mode="fan_out", nonlinearity="relu")640 # copied from the `reset_parameters` method of `class Linear(Module)` in `torch`.641 elif isinstance(module, nn.Linear):642 nn.init.kaiming_uniform_(module.weight, a=math.sqrt(5))643 if module.bias is not None:644 fan_in, _ = nn.init._calculate_fan_in_and_fan_out(module.weight)645 bound = 1 / math.sqrt(fan_in) if fan_in > 0 else 0646 nn.init.uniform_(module.bias, -bound, bound)647 elif isinstance(module, (nn.BatchNorm2d, nn.GroupNorm)):648 nn.init.constant_(module.weight, 1)649 nn.init.constant_(module.bias, 0)650 651 652@auto_docstring653class BitModel(BitPreTrainedModel):654 def __init__(self, config):655 super().__init__(config)656 self.config = config657 658 self.embedder = BitEmbeddings(config)659 660 self.encoder = BitEncoder(config)661 self.norm = (662 BitGroupNormActivation(config, num_channels=config.hidden_sizes[-1])663 if config.layer_type == "preactivation"664 else nn.Identity()665 )666 667 self.pooler = nn.AdaptiveAvgPool2d((1, 1))668 # Initialize weights and apply final processing669 self.post_init()670 671 @auto_docstring672 def forward(673 self, pixel_values: Tensor, output_hidden_states: Optional[bool] = None, return_dict: Optional[bool] = None674 ) -> BaseModelOutputWithPoolingAndNoAttention:675 output_hidden_states = (676 output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states677 )678 return_dict = return_dict if return_dict is not None else self.config.use_return_dict679 680 embedding_output = self.embedder(pixel_values)681 682 encoder_outputs = self.encoder(683 embedding_output, output_hidden_states=output_hidden_states, return_dict=return_dict684 )685 686 last_hidden_state = encoder_outputs[0]687 688 last_hidden_state = self.norm(last_hidden_state)689 690 pooled_output = self.pooler(last_hidden_state)691 692 if not return_dict:693 return (last_hidden_state, pooled_output) + encoder_outputs[1:]694 695 return BaseModelOutputWithPoolingAndNoAttention(696 last_hidden_state=last_hidden_state,697 pooler_output=pooled_output,698 hidden_states=encoder_outputs.hidden_states,699 )700 701 702@auto_docstring(703 custom_intro="""704 BiT Model with an image classification head on top (a linear layer on top of the pooled features), e.g. for705 ImageNet.706 """707)708class BitForImageClassification(BitPreTrainedModel):709 def __init__(self, config):710 super().__init__(config)711 self.num_labels = config.num_labels712 self.bit = BitModel(config)713 # classification head714 self.classifier = nn.Sequential(715 nn.Flatten(),716 nn.Linear(config.hidden_sizes[-1], config.num_labels) if config.num_labels > 0 else nn.Identity(),717 )718 # initialize weights and apply final processing719 self.post_init()720 721 @auto_docstring722 def forward(723 self,724 pixel_values: Optional[torch.FloatTensor] = None,725 labels: Optional[torch.LongTensor] = None,726 output_hidden_states: Optional[bool] = None,727 return_dict: Optional[bool] = None,728 ) -> ImageClassifierOutputWithNoAttention:729 r"""730 labels (`torch.LongTensor` of shape `(batch_size,)`, *optional*):731 Labels for computing the image classification/regression loss. Indices should be in `[0, ...,732 config.num_labels - 1]`. If `config.num_labels > 1` a classification loss is computed (Cross-Entropy).733 """734 return_dict = return_dict if return_dict is not None else self.config.use_return_dict735 736 outputs = self.bit(pixel_values, output_hidden_states=output_hidden_states, return_dict=return_dict)737 738 pooled_output = outputs.pooler_output if return_dict else outputs[1]739 740 logits = self.classifier(pooled_output)741 742 loss = None743 744 if labels is not None:745 loss = self.loss_function(labels, logits, self.config)746 747 if not return_dict:748 output = (logits,) + outputs[2:]749 return (loss,) + output if loss is not None else output750 751 return ImageClassifierOutputWithNoAttention(loss=loss, logits=logits, hidden_states=outputs.hidden_states)752 753 754@auto_docstring(755 custom_intro="""756 BiT backbone, to be used with frameworks like DETR and MaskFormer.757 """758)759class BitBackbone(BitPreTrainedModel, BackboneMixin):760 has_attentions = False761 762 def __init__(self, config):763 super().__init__(config)764 super()._init_backbone(config)765 766 self.bit = BitModel(config)767 self.num_features = [config.embedding_size] + config.hidden_sizes768 769 # initialize weights and apply final processing770 self.post_init()771 772 @auto_docstring773 def forward(774 self, pixel_values: Tensor, output_hidden_states: Optional[bool] = None, return_dict: Optional[bool] = None775 ) -> BackboneOutput:776 r"""777 Examples:778 779 ```python780 >>> from transformers import AutoImageProcessor, AutoBackbone781 >>> import torch782 >>> from PIL import Image783 >>> import requests784 785 >>> url = "http://images.cocodataset.org/val2017/000000039769.jpg"786 >>> image = Image.open(requests.get(url, stream=True).raw)787 788 >>> processor = AutoImageProcessor.from_pretrained("google/bit-50")789 >>> model = AutoBackbone.from_pretrained("google/bit-50")790 791 >>> inputs = processor(image, return_tensors="pt")792 >>> outputs = model(**inputs)793 ```"""794 return_dict = return_dict if return_dict is not None else self.config.use_return_dict795 output_hidden_states = (796 output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states797 )798 799 outputs = self.bit(pixel_values, output_hidden_states=True, return_dict=True)800 801 hidden_states = outputs.hidden_states802 803 feature_maps = ()804 for idx, stage in enumerate(self.stage_names):805 if stage in self.out_features:806 feature_maps += (hidden_states[idx],)807 808 if not return_dict:809 output = (feature_maps,)810 if output_hidden_states:811 output += (outputs.hidden_states,)812 return output813 814 return BackboneOutput(815 feature_maps=feature_maps,816 hidden_states=outputs.hidden_states if output_hidden_states else None,817 attentions=None,818 )819 820 821__all__ = ["BitForImageClassification", "BitModel", "BitPreTrainedModel", "BitBackbone"]822