diff --git a/README.md b/README.md index 70b78ecb..8b584ce2 100644 --- a/README.md +++ b/README.md @@ -119,6 +119,8 @@ The Markdown sources live in [`docs/`](docs/) if you would rather read them in t | NextViT | [Next-ViT: Next Generation Vision Transformer for Efficient Deployment in Realistic Industrial Scenarios](https://arxiv.org/abs/2207.05501) | `timm` | | PiT | [Rethinking Spatial Dimensions of Vision Transformers](https://arxiv.org/abs/2103.16302) | `timm` | | PoolFormer | [MetaFormer is Actually What You Need for Vision](https://arxiv.org/abs/2111.11418) | `timm` | + | PVT | [Pyramid Vision Transformer: A Versatile Backbone for Dense Prediction without Convolutions](https://arxiv.org/abs/2102.12122) | `transformers` | + | PVTv2 | [PVTv2: Improved Baselines with Pyramid Vision Transformer](https://arxiv.org/abs/2106.13797) | `transformers` | | Res2Net | [Res2Net: A New Multi-scale Backbone Architecture](https://arxiv.org/abs/1904.01169) | `timm` | | ResMLP | [ResMLP: Feedforward networks for image classification with data-efficient training](https://arxiv.org/abs/2105.03404) | `timm` | | ResNet | [Deep Residual Learning for Image Recognition](https://arxiv.org/abs/1512.03385) | `timm` | diff --git a/docs/pvt.md b/docs/pvt.md new file mode 100644 index 00000000..3f2d0ffa --- /dev/null +++ b/docs/pvt.md @@ -0,0 +1,160 @@ +# PVT + +
+Weights: pretrained Keras weights live on Hugging Face under +zeromodels/pvt-<variant>-224 +(each repo carries zm_config.json + model.weights.h5). +Load with from_weights("zeromodels/pvt-tiny-224"). +
+ +PVT (Pyramid Vision Transformer) is a hierarchical vision transformer: four stages that +halve the spatial resolution and grow the channel width, so a single backbone produces a +CNN-style feature pyramid usable for classification and dense prediction. Each stage is a +**non-overlapping** convolutional patch embedding with a **learned position embedding**, +**spatial-reduction attention** (the key/value tokens are shrunk by a strided convolution so +attention stays affordable at high resolution), and a standard two-dense feed-forward +network. The last stage prepends a class token, and the classifier reads it. + +**Paper**: [Pyramid Vision Transformer: A Versatile Backbone for Dense Prediction without Convolutions](https://arxiv.org/abs/2102.12122) + +For the second-generation model (overlapping patches, no position embeddings, convolutional +FFN, and a linear-attention option), see [PVTv2](pvt_v2.md). + +## API + +### PvtImageClassify + +```python +PvtImageClassify( + hidden_sizes=(64, 128, 320, 512), + depths=(2, 2, 2, 2), + num_attention_heads=(1, 2, 5, 8), + sr_ratios=(8, 4, 2, 1), + mlp_ratios=(8, 8, 4, 4), + image_size=224, + include_normalization=True, + normalization_mode="imagenet", + num_classes=1000, + classifier_activation="linear", + name="PvtImageClassify", +) +``` + +The classifier: the backbone plus a dense head over the last stage's class token. +`include_normalization=True` means the model takes **raw `[0, 255]` pixels** and applies +ImageNet mean/std internally, so there is no separate image processor to construct. + +**Parameters** + +- **hidden_sizes** / **depths** / **num_attention_heads** / **sr_ratios** / **mlp_ratios** (`tuple`): per-stage width, block count, heads, spatial-reduction ratio, and FFN expansion. The variants differ only in `depths`; `from_weights` fills these from the variant config. +- **image_size** (`int`, *optional*, defaults to `224`): resolution the model is built for. The learned position embeddings are interpolated to this grid (see [Variable Input Resolution](#variable-input-resolution)). +- **include_normalization** (`bool`, *optional*, defaults to `True`): bake ImageNet normalization into the graph. +- **num_classes** (`int`, *optional*, defaults to `1000`): classifier outputs. + +**Call** `model(pixel_values, training=False)`. **Returns** class logits of shape `(B, num_classes)`. + +### PvtModel + +The backbone alone. With `as_backbone=True` it returns the four stage feature maps +(the pyramid, class token dropped) instead of just the last one, for detection or +segmentation necks. + +```python +PvtModel(as_backbone=False, hidden_sizes=(64, 128, 320, 512), ..., include_normalization=True) +``` + +### PvtConfig + +Typed config (`model_type="pvt"`) holding the fields above; serialized into each Hub repo's +`zm_config.json`. + +## Model Variants + +For `PvtImageClassify.from_weights("zeromodels/")`. Every variant shares the widths +`(64, 128, 320, 512)` and differs only in depth: + +| Variant id | Depths | Params | ImageNet-1k top-1 | +|-------------------|---------------|-------:|------------------:| +| `pvt-tiny-224` | (2, 2, 2, 2) | 13.2M | 75.1% | +| `pvt-small-224` | (3, 4, 6, 3) | 24.5M | 79.8% | +| `pvt-medium-224` | (3, 4, 18, 3) | 44.2M | 81.2% | +| `pvt-large-224` | (3, 8, 27, 3) | 61.4M | 81.7% | + +Reported top-1 is from the paper. All variants are 224x224, 1000 classes. + +## Basic Usage + +```python +import keras +import numpy as np +from PIL import Image +from zeromodels.models.pvt import PvtImageClassify + +model = PvtImageClassify.from_weights("zeromodels/pvt-tiny-224") + +image = Image.open("assets/data/hf_cat_2.jpg").convert("RGB").resize((224, 224)) +pixels = np.asarray(image, "float32")[None] # (1, 224, 224, 3), raw [0, 255] + +logits = model(pixels, training=False) +top5 = np.argsort(keras.ops.convert_to_numpy(logits)[0])[-5:][::-1] +print("top-5 ImageNet-1k class ids:", top5.tolist()) +``` + +Normalization is inside the model, so pass raw pixels. Map the class ids to the +[ImageNet-1k label list](https://huggingface.co/datasets/imagenet-1k) to read names. + +## Feature Pyramid + +For detection / segmentation, take the four stage outputs: + +```python +from zeromodels.models.pvt import PvtModel + +backbone = PvtModel.from_weights("zeromodels/pvt-tiny-224", as_backbone=True) +feats = backbone(np.zeros((1, 224, 224, 3), "float32"), training=False) +print([tuple(f.shape) for f in feats]) +# [(1, 56, 56, 64), (1, 28, 28, 128), (1, 14, 14, 320), (1, 7, 7, 512)] +``` + +The strides are 4, 8, 16, 32, matching a standard CNN backbone. + +## Variable Input Resolution + +Unlike [PVTv2](pvt_v2.md), PVT v1 has **learned position embeddings**, so a non-224 input +needs them resized. Build the model at the target size and `from_weights` bilinearly +interpolates each stage's position embedding from its trained 224 grid to the new grid at +load time. + +```python +model = PvtImageClassify.from_weights("zeromodels/pvt-tiny-224", image_size=384) +logits = model(np.zeros((1, 384, 384, 3), "float32"), training=False) +``` + +## Data Format + +**The model supports both `channels_last` and `channels_first`, and the two are +bit-exact.** A model reads `keras.config.image_data_format()` when it is **constructed** +(there is no `data_format` argument); set the format before building. + +```python +import keras +keras.config.set_image_data_format("channels_first") +model = PvtImageClassify.from_weights("zeromodels/pvt-tiny-224") # expects (B, 3, H, W) +``` + +## Loading Fine-tuned and Community Weights + +Any Hugging Face repo whose `model_type` is `"pvt"` (for example the original +`Zetatech/pvt-*-224` checkpoints) loads with the `hf:` prefix, converting on the fly: + +```python +from zeromodels.models.pvt import PvtImageClassify + +model = PvtImageClassify.from_weights("hf:Zetatech/pvt-tiny-224") +model = PvtImageClassify.from_weights("hf:/pvt-finetuned-on-my-data") + +# Architecture only, randomly initialized +model = PvtImageClassify.from_weights("zeromodels/pvt-tiny-224", load_weights=False) +``` + +`PvtModel` accepts `hf:` the same way. diff --git a/docs/pvt_v2.md b/docs/pvt_v2.md new file mode 100644 index 00000000..da7a866e --- /dev/null +++ b/docs/pvt_v2.md @@ -0,0 +1,166 @@ +# PVTv2 + +
+Weights: pretrained Keras weights live on Hugging Face under +zeromodels/pvt-v2-<variant> +(each repo carries zm_config.json + model.weights.h5). +Load with from_weights("zeromodels/pvt-v2-b0"). +
+ +PVTv2 is a hierarchical vision transformer: four stages that halve the spatial resolution +and grow the channel width, so a single backbone produces a feature pyramid the way a CNN +does. It improves on [PVT](pvt.md) in three ways: an **overlapping** convolutional patch +embedding, **spatial-reduction attention** that shrinks the key/value sequence with a +strided convolution (so attention stays affordable at high resolution), and a +**convolutional feed-forward network** (a 3x3 depthwise conv between the two dense layers) +that removes the need for any position embedding. Dropping position embeddings is what lets +it run at arbitrary input resolution with no interpolation. + +The `b2_linear` variant swaps spatial-reduction attention for **linear attention**: instead +of a strided conv, it average-pools every stage to a fixed 7x7 grid, so the key/value length +is constant regardless of input size and the cost is linear in the number of tokens. + +**Paper**: [PVTv2: Improved Baselines with Pyramid Vision Transformer](https://arxiv.org/abs/2106.13797) + +For the first-generation model (non-overlapping patches, learned position embeddings), see +[PVT](pvt.md). + +## API + +### PvtV2ImageClassify + +```python +PvtV2ImageClassify( + hidden_sizes=(32, 64, 160, 256), + depths=(2, 2, 2, 2), + num_attention_heads=(1, 2, 5, 8), + sr_ratios=(8, 4, 2, 1), + mlp_ratios=(8, 8, 4, 4), + linear_attention=False, + image_size=224, + include_normalization=True, + normalization_mode="imagenet", + num_classes=1000, + classifier_activation="linear", + name="PvtV2ImageClassify", +) +``` + +The classifier: the backbone, a global average pool over the last stage, and one dense head. +`include_normalization=True` means the model takes **raw `[0, 255]` pixels** and applies +ImageNet mean/std internally, so there is no separate image processor to construct. + +**Parameters** + +- **hidden_sizes** / **depths** / **num_attention_heads** / **sr_ratios** / **mlp_ratios** (`tuple`): per-stage width, block count, heads, spatial-reduction ratio, and FFN expansion. `from_weights` fills these from the variant config. +- **linear_attention** (`bool`, *optional*, defaults to `False`): use the fixed-7x7 pooled linear-attention variant (`b2_linear`). +- **image_size** (`int`, *optional*, defaults to `224`): resolution the model is built for. +- **include_normalization** (`bool`, *optional*, defaults to `True`): bake ImageNet normalization into the graph. +- **num_classes** (`int`, *optional*, defaults to `1000`): classifier outputs. + +**Call** `model(pixel_values, training=False)`. **Returns** class logits of shape `(B, num_classes)`. + +### PvtV2Model + +The backbone alone. With `as_backbone=True` it returns the four stage feature maps +(the pyramid) instead of just the last one, for detection or segmentation necks. + +```python +PvtV2Model(as_backbone=False, hidden_sizes=(32, 64, 160, 256), ..., include_normalization=True) +``` + +### PvtV2Config + +Typed config (`model_type="pvt_v2"`) holding the fields above; serialized into each Hub +repo's `zm_config.json`. + +## Model Variants + +For `PvtV2ImageClassify.from_weights("zeromodels/")`: + +| Variant id | Params | ImageNet-1k top-1 | Notes | +|--------------------|-------:|------------------:|---------------------------| +| `pvt-v2-b0` | 3.7M | 70.5% | | +| `pvt-v2-b1` | 14.0M | 78.7% | | +| `pvt-v2-b2` | 25.4M | 82.0% | | +| `pvt-v2-b2-linear` | 22.6M | 82.1% | linear (pooled) attention | +| `pvt-v2-b3` | 45.2M | 83.1% | | +| `pvt-v2-b4` | 62.6M | 83.6% | | +| `pvt-v2-b5` | 82.0M | 83.8% | | + +Reported top-1 is from the paper. All variants are 224x224, 1000 classes. + +## Basic Usage + +```python +import keras +import numpy as np +from PIL import Image +from zeromodels.models.pvt_v2 import PvtV2ImageClassify + +model = PvtV2ImageClassify.from_weights("zeromodels/pvt-v2-b2") + +image = Image.open("assets/data/hf_cat_2.jpg").convert("RGB").resize((224, 224)) +pixels = np.asarray(image, "float32")[None] # (1, 224, 224, 3), raw [0, 255] + +logits = model(pixels, training=False) +top5 = np.argsort(keras.ops.convert_to_numpy(logits)[0])[-5:][::-1] +print("top-5 ImageNet-1k class ids:", top5.tolist()) +``` + +Normalization is inside the model, so pass raw pixels. Map the class ids to the +[ImageNet-1k label list](https://huggingface.co/datasets/imagenet-1k) to read names. + +## Feature Pyramid + +For detection / segmentation, take the four stage outputs: + +```python +from zeromodels.models.pvt_v2 import PvtV2Model + +backbone = PvtV2Model.from_weights("zeromodels/pvt-v2-b2", as_backbone=True) +feats = backbone(np.zeros((1, 224, 224, 3), "float32"), training=False) +print([tuple(f.shape) for f in feats]) +# [(1, 56, 56, 64), (1, 28, 28, 128), (1, 14, 14, 320), (1, 7, 7, 512)] +``` + +The strides are 4, 8, 16, 32, matching a standard CNN backbone. + +## Variable Input Resolution + +PVTv2 has **no position embeddings**, so any resolution works with no interpolation: build +the model at the size you want. + +```python +model = PvtV2ImageClassify.from_weights("zeromodels/pvt-v2-b2", image_size=384) +logits = model(np.zeros((1, 384, 384, 3), "float32"), training=False) +``` + +## Data Format + +**The model supports both `channels_last` and `channels_first`, and the two are +bit-exact.** A model reads `keras.config.image_data_format() `when it is **constructed** +(there is no `data_format` argument); set the format before building. + +```python +import keras +keras.config.set_image_data_format("channels_first") +model = PvtV2ImageClassify.from_weights("zeromodels/pvt-v2-b0") # expects (B, 3, H, W) +``` + +## Loading Fine-tuned and Community Weights + +Any Hugging Face repo whose `model_type` is `"pvt_v2"` (for example the original +`OpenGVLab/pvt_v2_*` checkpoints) loads with the `hf:` prefix, converting on the fly: + +```python +from zeromodels.models.pvt_v2 import PvtV2ImageClassify + +model = PvtV2ImageClassify.from_weights("hf:OpenGVLab/pvt_v2_b2") +model = PvtV2ImageClassify.from_weights("hf:/pvt-v2-finetuned-on-my-data") + +# Architecture only, randomly initialized +model = PvtV2ImageClassify.from_weights("zeromodels/pvt-v2-b2", load_weights=False) +``` + +`PvtV2Model` accepts `hf:` the same way. diff --git a/website/mkdocs.yml b/website/mkdocs.yml index 8bf1fd02..5d449ba3 100644 --- a/website/mkdocs.yml +++ b/website/mkdocs.yml @@ -182,6 +182,8 @@ nav: - MaskFormer: maskformer.md - MobileViT: mobilevit.md - MobileViTV2: mobilevitv2.md + - PVT: pvt.md + - PVTv2: pvt_v2.md - RF-DETR: rf_detr.md - RT-DETR: rt_detr.md - RT-DETRv2: rt_detr_v2.md diff --git a/zeromodels/conversion/weight_transfer_util.py b/zeromodels/conversion/weight_transfer_util.py index c1ef9365..7aeebece 100644 --- a/zeromodels/conversion/weight_transfer_util.py +++ b/zeromodels/conversion/weight_transfer_util.py @@ -204,7 +204,15 @@ def transform_conv_weights( elif any( substring in keras_name.lower() - for substring in ["conv", "conv2d", "pointwise", "downsample", "sr"] + for substring in [ + "conv", + "conv2d", + "pointwise", + "downsample", + "sr", + "reduction", # spatial_reduction / sequence_reduction (PVT SRA conv) + "proj", # patch-embed proj / projection conv (only reached for 4D weights) + ] ): # Standard 2D convolution return np.transpose(torch_weight, [2, 3, 1, 0]) diff --git a/zeromodels/models/__init__.py b/zeromodels/models/__init__.py index a413d52c..16eeb55c 100644 --- a/zeromodels/models/__init__.py +++ b/zeromodels/models/__init__.py @@ -87,6 +87,8 @@ owlvit, pit, poolformer, + pvt, + pvt_v2, qwen2, qwen2_5_vl, qwen2_moe, diff --git a/zeromodels/models/pvt/__init__.py b/zeromodels/models/pvt/__init__.py new file mode 100644 index 00000000..8bd6f9d9 --- /dev/null +++ b/zeromodels/models/pvt/__init__.py @@ -0,0 +1,4 @@ +from zeromodels.models.pvt.pvt_config import PVT_VARIANTS, PvtConfig +from zeromodels.models.pvt.pvt_model import PvtImageClassify, PvtModel + +__all__ = ["PvtImageClassify", "PvtModel", "PvtConfig", "PVT_VARIANTS"] diff --git a/zeromodels/models/pvt/convert_pvt_hf_to_keras.py b/zeromodels/models/pvt/convert_pvt_hf_to_keras.py new file mode 100644 index 00000000..da9d4f36 --- /dev/null +++ b/zeromodels/models/pvt/convert_pvt_hf_to_keras.py @@ -0,0 +1,165 @@ +import gc +import re + +import keras +from tqdm import tqdm + +from zeromodels.conversion import verify_cls_model_equivalence +from zeromodels.conversion.exceptions import ( + WeightMappingError, + WeightShapeMismatchError, +) +from zeromodels.conversion.hf_download_utils import download_hf_state_dict +from zeromodels.conversion.weight_split_util import split_model_weights +from zeromodels.conversion.weight_transfer_util import ( + compare_keras_torch_names, + transfer_attention_weights, + transfer_weights, +) +from zeromodels.models.pvt import PvtImageClassify +from zeromodels.models.pvt.pvt_config import PVT_VARIANTS + +PVT_MODEL_CONFIG = { + "pvt_tiny": { + "hidden_sizes": (64, 128, 320, 512), + "depths": (2, 2, 2, 2), + "num_attention_heads": (1, 2, 5, 8), + "sr_ratios": (8, 4, 2, 1), + "mlp_ratios": (8, 8, 4, 4), + "image_size": 224, + "num_classes": 1000, + }, + "pvt_small": { + "hidden_sizes": (64, 128, 320, 512), + "depths": (3, 4, 6, 3), + "num_attention_heads": (1, 2, 5, 8), + "sr_ratios": (8, 4, 2, 1), + "mlp_ratios": (8, 8, 4, 4), + "image_size": 224, + "num_classes": 1000, + }, + "pvt_medium": { + "hidden_sizes": (64, 128, 320, 512), + "depths": (3, 4, 18, 3), + "num_attention_heads": (1, 2, 5, 8), + "sr_ratios": (8, 4, 2, 1), + "mlp_ratios": (8, 8, 4, 4), + "image_size": 224, + "num_classes": 1000, + }, + "pvt_large": { + "hidden_sizes": (64, 128, 320, 512), + "depths": (3, 8, 27, 3), + "num_attention_heads": (1, 2, 5, 8), + "sr_ratios": (8, 4, 2, 1), + "mlp_ratios": (8, 8, 4, 4), + "image_size": 224, + "num_classes": 1000, + }, +} + +WEIGHT_NAME_MAPPING = { + "_": ".", + "patch.embed": "pvt.encoder.patch_embeddings", + "final.layernorm": "pvt.encoder.layer_norm", + "block": "pvt.encoder.block", + "layernorm.1": "layer_norm_1", + "layernorm.2": "layer_norm_2", + "layernorm": "layer_norm", + "proj": "projection", + "kernel": "weight", + "gamma": "weight", + "beta": "bias", + "predictions": "classifier", +} + +ATTENTION_NAME_MAPPING = { + "block": "pvt.encoder.block", + "attn.query": "attention.self.query", + "attn.key": "attention.self.key", + "attn.value": "attention.self.value", + "attn.proj": "attention.output.dense", + "attn.sr": "attention.self.sequence_reduction", + "attn.norm": "attention.self.layer_norm", +} + + +def transfer_pvt_weights(keras_model, state_dict): + trainable, non_trainable = split_model_weights(keras_model) + + for keras_weight, keras_name in tqdm( + trainable + non_trainable, desc="Transferring weights to Keras" + ): + leaf = keras_weight.path.split("/")[-1] + if leaf in ("cls_token", "pos_embed"): + i = re.search(r"patch_embed_(\d+)_", keras_weight.path).group(1) + comp = "cls_token" if leaf == "cls_token" else "position_embeddings" + torch_name = f"pvt.encoder.patch_embeddings.{i}.{comp}" + keras_weight.assign(state_dict[torch_name]) + continue + + torch_name = keras_name + for old, new in WEIGHT_NAME_MAPPING.items(): + torch_name = torch_name.replace(old, new) + + if "attention" in torch_name: + transfer_attention_weights( + keras_name, keras_weight, state_dict, ATTENTION_NAME_MAPPING + ) + continue + + if torch_name not in state_dict: + raise WeightMappingError(keras_name, torch_name) + + torch_weight = state_dict[torch_name] + if not compare_keras_torch_names( + keras_name, keras_weight, torch_name, torch_weight + ): + raise WeightShapeMismatchError( + keras_name, keras_weight.shape, torch_name, torch_weight.shape + ) + transfer_weights(keras_name, keras_weight, torch_weight) + + +if __name__ == "__main__": + import importlib.metadata as _meta + + _orig_version = _meta.version + _meta.version = lambda name: "0.23.0" if name == "tokenizers" else _orig_version(name) + import transformers + + for variant, meta in PVT_VARIANTS.items(): + hf_id = meta["hf_id"] + print(f"\n{'=' * 60}") + print(f"Converting: {variant} <- {hf_id}") + print(f"{'=' * 60}") + + state = download_hf_state_dict(hf_id) + keras_model = PvtImageClassify( + **PVT_MODEL_CONFIG[meta["model"]], include_normalization=False + ) + transfer_pvt_weights(keras_model, state) + + hf_model = transformers.PvtForImageClassification.from_pretrained(hf_id).eval() + results = verify_cls_model_equivalence( + model_a=hf_model, + model_b=keras_model, + input_shape=keras_model.input_shape[1:], + output_specs={"num_classes": keras_model.output_shape[-1]}, + comparison_type="hf_to_keras", + run_performance=False, + atol=1e-2, + rtol=1e-2, + ) + if not results["standard_input"]: + raise ValueError( + "Model equivalence test failed - outputs do not match for standard input" + ) + + out_path = f"{variant}.weights.h5" + keras_model.save_weights(out_path) + print(f" Saved -> {out_path}") + + del keras_model, state, hf_model + keras.backend.clear_session() + gc.collect() diff --git a/zeromodels/models/pvt/pvt_config.py b/zeromodels/models/pvt/pvt_config.py new file mode 100644 index 00000000..51dcb57a --- /dev/null +++ b/zeromodels/models/pvt/pvt_config.py @@ -0,0 +1,59 @@ +from zeromodels.base import BaseConfig + + +class PvtConfig(BaseConfig): + r"""Configuration for [`PvtModel`] / [`PvtImageClassify`]. + + PVT (Pyramid Vision Transformer v1) is a hierarchical transformer: four stages, each a + non-overlapping convolutional patch embedding with a learned position embedding, + spatial-reduction attention (the key/value tokens are reduced by a strided conv), and a + standard two-Dense feed-forward network. The last stage prepends a class token, and the + classifier reads it. Variable input resolution is handled by bilinearly interpolating + each stage's position embedding at weight-load time. One `zm_config.json` (declaring the + canonical [`PvtImageClassify`]) sits on each variant's repo; both the backbone and the + classifier load from it. Fields mirror the model constructor and serialize flat. + + Args: + hidden_sizes (`tuple`, *optional*, defaults to `(64, 128, 320, 512)`): + Channel width per stage. + depths (`tuple`, *optional*, defaults to `(2, 2, 2, 2)`): + Number of transformer blocks per stage. + num_attention_heads (`tuple`, *optional*, defaults to `(1, 2, 5, 8)`): + Attention heads per stage. + sr_ratios (`tuple`, *optional*, defaults to `(8, 4, 2, 1)`): + Spatial-reduction ratio of the key/value tokens per stage. + mlp_ratios (`tuple`, *optional*, defaults to `(8, 8, 4, 4)`): + Feed-forward hidden expansion per stage. + image_size (`int`, *optional*, defaults to 224): + Square input resolution the weights were trained at. + num_classes (`int`, *optional*, defaults to 1000): + Number of classifier output classes (backbone ignores it). + + Examples: + + ```python + >>> from zeromodels.models.pvt import PvtConfig, PvtImageClassify + + >>> configuration = PvtConfig() + >>> model = PvtImageClassify(configuration) + >>> configuration = model.config + ```""" + + model_type = "pvt" + + hidden_sizes: tuple = (64, 128, 320, 512) + depths: tuple = (2, 2, 2, 2) + num_attention_heads: tuple = (1, 2, 5, 8) + sr_ratios: tuple = (8, 4, 2, 1) + mlp_ratios: tuple = (8, 8, 4, 4) + image_size: int = 224 + num_classes: int = 1000 + + +# Hosted variants -> arch preset. Weights load by Hub repo id (zm_config.json). +PVT_VARIANTS = { + "pvt_tiny": {"model": "pvt_tiny", "hf_id": "Zetatech/pvt-tiny-224"}, + "pvt_small": {"model": "pvt_small", "hf_id": "Zetatech/pvt-small-224"}, + "pvt_medium": {"model": "pvt_medium", "hf_id": "Zetatech/pvt-medium-224"}, + "pvt_large": {"model": "pvt_large", "hf_id": "Zetatech/pvt-large-224"}, +} diff --git a/zeromodels/models/pvt/pvt_layers.py b/zeromodels/models/pvt/pvt_layers.py new file mode 100644 index 00000000..428e6bbb --- /dev/null +++ b/zeromodels/models/pvt/pvt_layers.py @@ -0,0 +1,216 @@ +import keras +from keras import layers, ops + +from zeromodels.base.base_attention import fused_attention + + +@keras.saving.register_keras_serializable(package="zeromodels") +class PvtSelfAttention(layers.Layer): + """PVT spatial-reduction attention. Query is projected from the full token sequence; + keys/values come from a sequence reduced by a strided conv + LayerNorm (skipped when + ``sr_ratio == 1``). ``proj`` is the attention-output Dense.""" + + def __init__( + self, + hidden_size, + num_heads, + sr_ratio, + qkv_bias=True, + layer_norm_eps=1e-6, + block_prefix="block", + **kwargs, + ): + super().__init__(**kwargs) + assert hidden_size % num_heads == 0 + self.hidden_size = hidden_size + self.num_heads = num_heads + self.head_dim = hidden_size // num_heads + self.scale = self.head_dim**-0.5 + self.sr_ratio = sr_ratio + self.qkv_bias = qkv_bias + self.layer_norm_eps = layer_norm_eps + self.block_prefix = block_prefix + self.data_format = keras.config.image_data_format() + + self.query = layers.Dense( + hidden_size, use_bias=qkv_bias, name=f"{block_prefix}_attn_query" + ) + self.key = layers.Dense( + hidden_size, use_bias=qkv_bias, name=f"{block_prefix}_attn_key" + ) + self.value = layers.Dense( + hidden_size, use_bias=qkv_bias, name=f"{block_prefix}_attn_value" + ) + self.proj = layers.Dense(hidden_size, name=f"{block_prefix}_attn_proj") + if sr_ratio > 1: + self.sr = layers.Conv2D( + hidden_size, + sr_ratio, + strides=sr_ratio, + padding="valid", + data_format=self.data_format, + name=f"{block_prefix}_attn_sr", + ) + self.norm = layers.LayerNormalization( + axis=-1, epsilon=layer_norm_eps, name=f"{block_prefix}_attn_norm" + ) + + def split_heads(self, x): + b = ops.shape(x)[0] + return ops.transpose( + ops.reshape(x, (b, -1, self.num_heads, self.head_dim)), (0, 2, 1, 3) + ) + + def call(self, x, height, width, training=None): + q = self.split_heads(self.query(x)) + if self.sr_ratio > 1: + b = ops.shape(x)[0] + grid = ops.reshape(x, (b, height, width, self.hidden_size)) + if self.data_format == "channels_first": + grid = ops.transpose(grid, (0, 3, 1, 2)) + grid = self.sr(grid) + if self.data_format == "channels_first": + grid = ops.transpose(grid, (0, 2, 3, 1)) + kv_in = self.norm(ops.reshape(grid, (b, -1, self.hidden_size))) + else: + kv_in = x + k = self.split_heads(self.key(kv_in)) + v = self.split_heads(self.value(kv_in)) + + out = fused_attention(q, k, v, self.scale, training=training) + out = ops.transpose(out, (0, 2, 1, 3)) + out = ops.reshape(out, (ops.shape(x)[0], ops.shape(x)[1], self.hidden_size)) + return self.proj(out) + + def compute_output_shape(self, input_shape): + return (input_shape[0], input_shape[1], self.hidden_size) + + def get_config(self): + config = super().get_config() + config.update( + { + "hidden_size": self.hidden_size, + "num_heads": self.num_heads, + "sr_ratio": self.sr_ratio, + "qkv_bias": self.qkv_bias, + "layer_norm_eps": self.layer_norm_eps, + "block_prefix": self.block_prefix, + } + ) + return config + + +@keras.saving.register_keras_serializable(package="zeromodels") +class PvtClsToken(layers.Layer): + """Prepends a learnable class token (last stage only).""" + + def build(self, input_shape): + self.cls_token = self.add_weight( + name="cls_token", + shape=(1, 1, input_shape[-1]), + initializer="zeros", + trainable=True, + ) + self.built = True + + def call(self, x): + cls = ops.broadcast_to(self.cls_token, (ops.shape(x)[0], 1, ops.shape(x)[-1])) + return ops.concatenate([cls, x], axis=1) + + def compute_output_shape(self, input_shape): + n = None if input_shape[1] is None else input_shape[1] + 1 + return (input_shape[0], n, input_shape[2]) + + +@keras.saving.register_keras_serializable(package="zeromodels") +class PvtStagePositions(layers.Layer): + """Adds a per-stage learned position embedding sized to ``grid_h x grid_w`` (plus a + class-token slot when ``has_cls``). On weight load from a different grid, the spatial + part is bilinearly interpolated to this model's grid (the class slot is kept).""" + + def __init__( + self, grid_h, grid_w, has_cls=False, resize_mode="bilinear", name=None, **kwargs + ): + super().__init__(name=name, **kwargs) + self.grid_h = int(grid_h) + self.grid_w = int(grid_w) + self.has_cls = has_cls + self.resize_mode = resize_mode + + def build(self, input_shape): + n = self.grid_h * self.grid_w + (1 if self.has_cls else 0) + self.pos_embed = self.add_weight( + name="pos_embed", + shape=(1, n, input_shape[-1]), + initializer="random_normal", + trainable=True, + ) + self.built = True + + def call(self, x): + return x + self.pos_embed + + def compute_output_shape(self, input_shape): + return input_shape + + def save_own_variables(self, store): + super().save_own_variables(store) + store["grid_h"] = self.grid_h + store["grid_w"] = self.grid_w + + def load_own_variables(self, store): + source_h, source_w = int(store["grid_h"][...]), int(store["grid_w"][...]) + if source_h == self.grid_h and source_w == self.grid_w: + self.pos_embed.assign(store["0"]) + return + pe = store["0"] + cls_pe, spatial = (pe[:, :1], pe[:, 1:]) if self.has_cls else (None, pe) + c = spatial.shape[-1] + spatial = ops.reshape(ops.cast(spatial, "float32"), (1, source_h, source_w, c)) + spatial = ops.image.resize( + spatial, + (self.grid_h, self.grid_w), + interpolation=self.resize_mode, + antialias=True, + ) + spatial = ops.reshape(spatial, (1, self.grid_h * self.grid_w, c)) + pe = ops.concatenate([cls_pe, spatial], axis=1) if self.has_cls else spatial + self.pos_embed.assign(pe) + + def get_config(self): + config = super().get_config() + config.update( + { + "grid_h": self.grid_h, + "grid_w": self.grid_w, + "has_cls": self.has_cls, + "resize_mode": self.resize_mode, + } + ) + return config + + +@keras.saving.register_keras_serializable(package="zeromodels") +class PvtDropPath(layers.Layer): + """Stochastic depth (identity at inference).""" + + def __init__(self, drop_prob, seed=None, **kwargs): + super().__init__(**kwargs) + self.drop_prob = drop_prob + self.seed = seed + self.seed_generator = keras.random.SeedGenerator(seed) + + def call(self, x, training=None): + if training and self.drop_prob > 0: + keep = 1 - self.drop_prob + shape = (ops.shape(x)[0],) + (1,) * (len(x.shape) - 1) + mask = ops.floor( + keep + keras.random.uniform(shape, 0, 1, seed=self.seed_generator) + ) + return (x / keep) * mask + return x + + def get_config(self): + config = super().get_config() + config.update({"drop_prob": self.drop_prob, "seed": self.seed}) + return config diff --git a/zeromodels/models/pvt/pvt_model.py b/zeromodels/models/pvt/pvt_model.py new file mode 100644 index 00000000..5fd20231 --- /dev/null +++ b/zeromodels/models/pvt/pvt_model.py @@ -0,0 +1,365 @@ +import keras +from keras import layers, ops + +from zeromodels.base import BaseModel +from zeromodels.conversion import copy_weights_by_path_suffix +from zeromodels.utils import standardize_input_shape +from zeromodels.utils.image_util import normalize_image_for_classify_models + +from .pvt_config import PvtConfig +from .pvt_layers import PvtClsToken, PvtDropPath, PvtSelfAttention, PvtStagePositions + +PVT_HUB_SIBLINGS = frozenset({"PvtModel", "PvtImageClassify"}) +PATCH_SIZES = (4, 2, 2, 2) +STRIDES = (4, 2, 2, 2) + + +def grid_to_tokens(x, channels, data_format): + if data_format == "channels_first": + x = ops.transpose(x, (0, 2, 3, 1)) + return layers.Reshape((-1, channels))(x) + + +def tokens_to_grid(x, H, W, channels, data_format): + x = layers.Reshape((H, W, channels))(x) + if data_format == "channels_first": + x = ops.transpose(x, (0, 3, 1, 2)) + return x + + +def pvt_mlp(x, channels, mid, name_prefix): + x = layers.Dense(mid, name=f"{name_prefix}_dense1")(x) + x = layers.Activation("gelu")(x) + x = layers.Dense(channels, name=f"{name_prefix}_dense2")(x) + return x + + +def pvt_block( + x, H, W, dim, num_heads, sr_ratio, mlp_ratio, drop_prob, stage_idx, block_idx +): + prefix = f"block_{stage_idx}_{block_idx}" + drop_path = PvtDropPath(drop_prob) + norm1 = layers.LayerNormalization( + axis=-1, epsilon=1e-6, name=f"{prefix}_layernorm_1" + )(x) + attn = PvtSelfAttention(dim, num_heads, sr_ratio, block_prefix=prefix)( + norm1, height=H, width=W + ) + x = layers.Add()([x, drop_path(attn)]) + norm2 = layers.LayerNormalization( + axis=-1, epsilon=1e-6, name=f"{prefix}_layernorm_2" + )(x) + mlp = pvt_mlp(norm2, dim, int(dim * mlp_ratio), name_prefix=f"{prefix}_mlp") + return layers.Add()([x, drop_path(mlp)]) + + +def pvt_backbone_feature( + inputs, + *, + hidden_sizes, + depths, + num_attention_heads, + sr_ratios, + mlp_ratios, + drop_path_rate, + data_format, + return_stages=False, +): + total = sum(depths) + dpr = [drop_path_rate * i / max(total - 1, 1) for i in range(total)] + x = inputs + features = [] + cur = 0 + for i in range(4): + x = layers.Conv2D( + hidden_sizes[i], + PATCH_SIZES[i], + strides=STRIDES[i], + padding="valid", + data_format=data_format, + name=f"patch_embed_{i}_proj", + )(x) + if data_format == "channels_first": + H, W = int(x.shape[2]), int(x.shape[3]) + else: + H, W = int(x.shape[1]), int(x.shape[2]) + x = grid_to_tokens(x, hidden_sizes[i], data_format) + x = layers.LayerNormalization( + axis=-1, epsilon=1e-6, name=f"patch_embed_{i}_layernorm" + )(x) + has_cls = i == 3 + if has_cls: + x = PvtClsToken(name=f"patch_embed_{i}_cls")(x) + x = PvtStagePositions(H, W, has_cls=has_cls, name=f"patch_embed_{i}_pos")(x) + for j in range(depths[i]): + x = pvt_block( + x, + H, + W, + hidden_sizes[i], + num_attention_heads[i], + sr_ratios[i], + mlp_ratios[i], + dpr[cur], + i, + j, + ) + cur += 1 + if i != 3: + x = tokens_to_grid(x, H, W, hidden_sizes[i], data_format) + features.append(x) + else: + x = layers.LayerNormalization( + axis=-1, epsilon=1e-6, name="final_layernorm" + )(x) + if return_stages: + patches = layers.Lambda(lambda v: v[:, 1:], name="drop_cls")(x) + features.append( + tokens_to_grid(patches, H, W, hidden_sizes[i], data_format) + ) + return features if return_stages else x + + +@keras.saving.register_keras_serializable(package="zeromodels") +class PvtModel(BaseModel): + """Pyramid Vision Transformer (PVT v1) backbone. + + Four hierarchical stages, each a non-overlapping convolutional patch embedding with a + learned position embedding, spatial-reduction attention, and a standard feed-forward + network; the last stage prepends a class token. The default output is the final + (class-token-carrying) token sequence used by the classifier; ``as_backbone=True`` + returns the four per-stage spatial feature maps. Variable input resolution is supported + by interpolating each stage's position embedding on weight load. + + References: + - [Pyramid Vision Transformer](https://arxiv.org/abs/2102.12122) + + Args: + See :class:`PvtConfig`. ``include_normalization`` bakes ImageNet normalization into + the graph; ``image_size`` sets the input the model is built for. Defaults describe + PVT-Tiny. + """ + + BASE_WEIGHT_CONFIG = None + config_class = PvtConfig + HUB_REPO_SIBLINGS = PVT_HUB_SIBLINGS + HF_MODEL_TYPE = "pvt" + + @classmethod + def from_hub_repo(cls, repo_id, load_weights=True, skip_mismatch=False, **kwargs): + model = cls.build_from_hub_repo(repo_id, **kwargs) + if load_weights: + src = PvtImageClassify.from_weights(repo_id, skip_mismatch=skip_mismatch) + copy_weights_by_path_suffix(src, model) + del src + return model + + @classmethod + def config_from_hf(cls, hf_config): + return { + "hidden_sizes": hf_config["hidden_sizes"], + "depths": hf_config["depths"], + "num_attention_heads": hf_config["num_attention_heads"], + "sr_ratios": hf_config["sequence_reduction_ratios"], + "mlp_ratios": hf_config["mlp_ratios"], + } + + @classmethod + def transfer_from_hf(cls, keras_model, state_dict): + from .convert_pvt_hf_to_keras import transfer_pvt_weights + + transfer_pvt_weights(keras_model, state_dict) + + def __init__( + self, + as_backbone=False, + hidden_sizes=(64, 128, 320, 512), + depths=(2, 2, 2, 2), + num_attention_heads=(1, 2, 5, 8), + sr_ratios=(8, 4, 2, 1), + mlp_ratios=(8, 8, 4, 4), + drop_path_rate=0.0, + image_size=224, + include_normalization=True, + normalization_mode="imagenet", + input_tensor=None, + name="PvtModel", + **kwargs, + ): + for k in ("num_classes", "classifier_activation", "hf_id"): + kwargs.pop(k, None) + + data_format = keras.config.image_data_format() + image_size = standardize_input_shape(image_size, data_format) + + if input_tensor is None: + img_input = layers.Input(shape=image_size) + elif not keras.utils.is_keras_tensor(input_tensor): + img_input = layers.Input(tensor=input_tensor, shape=image_size) + else: + img_input = input_tensor + + x = ( + normalize_image_for_classify_models(img_input, normalization_mode) + if include_normalization + else img_input + ) + features = pvt_backbone_feature( + x, + hidden_sizes=hidden_sizes, + depths=depths, + num_attention_heads=num_attention_heads, + sr_ratios=sr_ratios, + mlp_ratios=mlp_ratios, + drop_path_rate=drop_path_rate, + data_format=data_format, + return_stages=as_backbone, + ) + super().__init__(inputs=img_input, outputs=features, name=name, **kwargs) + + self.as_backbone = as_backbone + self.hidden_sizes = list(hidden_sizes) + self.depths = list(depths) + self.num_attention_heads = list(num_attention_heads) + self.sr_ratios = list(sr_ratios) + self.mlp_ratios = list(mlp_ratios) + self.drop_path_rate = drop_path_rate + self.image_size = image_size + self.include_normalization = include_normalization + self.normalization_mode = normalization_mode + self.input_tensor = input_tensor + + def get_config(self): + config = super().get_config() + config.update( + { + "as_backbone": self.as_backbone, + "hidden_sizes": self.hidden_sizes, + "depths": self.depths, + "num_attention_heads": self.num_attention_heads, + "sr_ratios": self.sr_ratios, + "mlp_ratios": self.mlp_ratios, + "drop_path_rate": self.drop_path_rate, + "image_size": self.image_size, + "include_normalization": self.include_normalization, + "normalization_mode": self.normalization_mode, + "input_tensor": self.input_tensor, + "name": self.name, + } + ) + return config + + @classmethod + def from_config(cls, config): + return cls(**config) + + +@keras.saving.register_keras_serializable(package="zeromodels") +class PvtImageClassify(BaseModel): + """PVT v1 image classifier: :class:`PvtModel` backbone + a Dense head on the class token. + + References: + - [Pyramid Vision Transformer](https://arxiv.org/abs/2102.12122) + + Args: + See :class:`PvtConfig`. ``num_classes`` / ``classifier_activation`` are + head-specific; all other args forward to :class:`PvtModel`. + """ + + BASE_WEIGHT_CONFIG = None + config_class = PvtConfig + HUB_REPO_SIBLINGS = PVT_HUB_SIBLINGS + HF_MODEL_TYPE = "pvt" + + @classmethod + def config_from_hf(cls, hf_config): + return { + "hidden_sizes": hf_config["hidden_sizes"], + "depths": hf_config["depths"], + "num_attention_heads": hf_config["num_attention_heads"], + "sr_ratios": hf_config["sequence_reduction_ratios"], + "mlp_ratios": hf_config["mlp_ratios"], + "num_classes": hf_config.get("num_labels", 1000), + } + + @classmethod + def transfer_from_hf(cls, keras_model, state_dict): + from .convert_pvt_hf_to_keras import transfer_pvt_weights + + transfer_pvt_weights(keras_model, state_dict) + + def __init__( + self, + hidden_sizes=(64, 128, 320, 512), + depths=(2, 2, 2, 2), + num_attention_heads=(1, 2, 5, 8), + sr_ratios=(8, 4, 2, 1), + mlp_ratios=(8, 8, 4, 4), + drop_path_rate=0.0, + image_size=224, + include_normalization=True, + normalization_mode="imagenet", + input_tensor=None, + num_classes=1000, + classifier_activation="linear", + name="PvtImageClassify", + **kwargs, + ): + kwargs.pop("hf_id", None) + + backbone = PvtModel( + hidden_sizes=hidden_sizes, + depths=depths, + num_attention_heads=num_attention_heads, + sr_ratios=sr_ratios, + mlp_ratios=mlp_ratios, + drop_path_rate=drop_path_rate, + image_size=image_size, + include_normalization=include_normalization, + normalization_mode=normalization_mode, + input_tensor=input_tensor, + name=f"{name}_backbone", + ) + tok = layers.Lambda(lambda v: v[:, 0], name="ExtractClsToken")(backbone.output) + out = layers.Dense( + num_classes, activation=classifier_activation, name="predictions" + )(tok) + super().__init__(inputs=backbone.input, outputs=out, name=name, **kwargs) + + self.hidden_sizes = list(hidden_sizes) + self.depths = list(depths) + self.num_attention_heads = list(num_attention_heads) + self.sr_ratios = list(sr_ratios) + self.mlp_ratios = list(mlp_ratios) + self.drop_path_rate = drop_path_rate + self.image_size = backbone.image_size + self.include_normalization = include_normalization + self.normalization_mode = normalization_mode + self.input_tensor = input_tensor + self.num_classes = num_classes + self.classifier_activation = classifier_activation + + def get_config(self): + config = super().get_config() + config.update( + { + "hidden_sizes": self.hidden_sizes, + "depths": self.depths, + "num_attention_heads": self.num_attention_heads, + "sr_ratios": self.sr_ratios, + "mlp_ratios": self.mlp_ratios, + "drop_path_rate": self.drop_path_rate, + "image_size": self.image_size, + "include_normalization": self.include_normalization, + "normalization_mode": self.normalization_mode, + "input_tensor": self.input_tensor, + "num_classes": self.num_classes, + "classifier_activation": self.classifier_activation, + "name": self.name, + } + ) + return config + + @classmethod + def from_config(cls, config): + return cls(**config) diff --git a/zeromodels/models/pvt_v2/__init__.py b/zeromodels/models/pvt_v2/__init__.py new file mode 100644 index 00000000..f9b3a3a8 --- /dev/null +++ b/zeromodels/models/pvt_v2/__init__.py @@ -0,0 +1,4 @@ +from zeromodels.models.pvt_v2.pvt_v2_config import PVT_V2_VARIANTS, PvtV2Config +from zeromodels.models.pvt_v2.pvt_v2_model import PvtV2ImageClassify, PvtV2Model + +__all__ = ["PvtV2ImageClassify", "PvtV2Model", "PvtV2Config", "PVT_V2_VARIANTS"] diff --git a/zeromodels/models/pvt_v2/convert_pvt_v2_hf_to_keras.py b/zeromodels/models/pvt_v2/convert_pvt_v2_hf_to_keras.py new file mode 100644 index 00000000..0f521cb4 --- /dev/null +++ b/zeromodels/models/pvt_v2/convert_pvt_v2_hf_to_keras.py @@ -0,0 +1,189 @@ +import gc + +import keras +from tqdm import tqdm + +from zeromodels.conversion import verify_cls_model_equivalence +from zeromodels.conversion.exceptions import ( + WeightMappingError, + WeightShapeMismatchError, +) +from zeromodels.conversion.hf_download_utils import download_hf_state_dict +from zeromodels.conversion.weight_split_util import split_model_weights +from zeromodels.conversion.weight_transfer_util import ( + compare_keras_torch_names, + transfer_attention_weights, + transfer_weights, +) +from zeromodels.models.pvt_v2 import PvtV2ImageClassify +from zeromodels.models.pvt_v2.pvt_v2_config import PVT_V2_VARIANTS + +PVT_V2_MODEL_CONFIG = { + "pvt_v2_b0": { + "hidden_sizes": (32, 64, 160, 256), + "depths": (2, 2, 2, 2), + "num_attention_heads": (1, 2, 5, 8), + "sr_ratios": (8, 4, 2, 1), + "mlp_ratios": (8, 8, 4, 4), + "linear_attention": False, + "image_size": 224, + "num_classes": 1000, + }, + "pvt_v2_b1": { + "hidden_sizes": (64, 128, 320, 512), + "depths": (2, 2, 2, 2), + "num_attention_heads": (1, 2, 5, 8), + "sr_ratios": (8, 4, 2, 1), + "mlp_ratios": (8, 8, 4, 4), + "linear_attention": False, + "image_size": 224, + "num_classes": 1000, + }, + "pvt_v2_b2": { + "hidden_sizes": (64, 128, 320, 512), + "depths": (3, 4, 6, 3), + "num_attention_heads": (1, 2, 5, 8), + "sr_ratios": (8, 4, 2, 1), + "mlp_ratios": (8, 8, 4, 4), + "linear_attention": False, + "image_size": 224, + "num_classes": 1000, + }, + "pvt_v2_b2_linear": { + "hidden_sizes": (64, 128, 320, 512), + "depths": (3, 4, 6, 3), + "num_attention_heads": (1, 2, 5, 8), + "sr_ratios": (8, 4, 2, 1), + "mlp_ratios": (8, 8, 4, 4), + "linear_attention": True, + "image_size": 224, + "num_classes": 1000, + }, + "pvt_v2_b3": { + "hidden_sizes": (64, 128, 320, 512), + "depths": (3, 4, 18, 3), + "num_attention_heads": (1, 2, 5, 8), + "sr_ratios": (8, 4, 2, 1), + "mlp_ratios": (8, 8, 4, 4), + "linear_attention": False, + "image_size": 224, + "num_classes": 1000, + }, + "pvt_v2_b4": { + "hidden_sizes": (64, 128, 320, 512), + "depths": (3, 8, 27, 3), + "num_attention_heads": (1, 2, 5, 8), + "sr_ratios": (8, 4, 2, 1), + "mlp_ratios": (8, 8, 4, 4), + "linear_attention": False, + "image_size": 224, + "num_classes": 1000, + }, + "pvt_v2_b5": { + "hidden_sizes": (64, 128, 320, 512), + "depths": (3, 6, 40, 3), + "num_attention_heads": (1, 2, 5, 8), + "sr_ratios": (8, 4, 2, 1), + "mlp_ratios": (4, 4, 4, 4), + "linear_attention": False, + "image_size": 224, + "num_classes": 1000, + }, +} + +WEIGHT_NAME_MAPPING = { + "_": ".", + "layers": "pvt_v2.encoder.layers", + "patch.embed": "patch_embedding", + "mlp.dwconv": "mlp.dwconv.dwconv", + "layernorm.1": "layer_norm_1", + "layernorm.2": "layer_norm_2", + "layernorm": "layer_norm", + "kernel": "weight", + "gamma": "weight", + "beta": "bias", + "predictions": "classifier", +} + +ATTENTION_NAME_MAPPING = { + "layers": "pvt_v2.encoder.layers", + "attn.query": "attention.query", + "attn.key": "attention.key", + "attn.value": "attention.value", + "attn.proj": "attention.proj", + "attn.sr": "attention.spatial_reduction", + "attn.norm": "attention.layer_norm", +} + + +def transfer_pvt_v2_weights(keras_model, state_dict): + trainable, non_trainable = split_model_weights(keras_model) + + for keras_weight, keras_name in tqdm( + trainable + non_trainable, desc="Transferring weights to Keras" + ): + torch_name = keras_name + for old, new in WEIGHT_NAME_MAPPING.items(): + torch_name = torch_name.replace(old, new) + + if "attention" in torch_name: + transfer_attention_weights( + keras_name, keras_weight, state_dict, ATTENTION_NAME_MAPPING + ) + continue + + if torch_name not in state_dict: + raise WeightMappingError(keras_name, torch_name) + + torch_weight = state_dict[torch_name] + if not compare_keras_torch_names( + keras_name, keras_weight, torch_name, torch_weight + ): + raise WeightShapeMismatchError( + keras_name, keras_weight.shape, torch_name, torch_weight.shape + ) + transfer_weights(keras_name, keras_weight, torch_weight) + + +if __name__ == "__main__": + import importlib.metadata as _meta + + _orig_version = _meta.version + _meta.version = lambda name: "0.23.0" if name == "tokenizers" else _orig_version(name) + import transformers + + for variant, meta in PVT_V2_VARIANTS.items(): + hf_id = meta["hf_id"] + print(f"\n{'=' * 60}") + print(f"Converting: {variant} <- {hf_id}") + print(f"{'=' * 60}") + + state = download_hf_state_dict(hf_id) + keras_model = PvtV2ImageClassify( + **PVT_V2_MODEL_CONFIG[meta["model"]], include_normalization=False + ) + transfer_pvt_v2_weights(keras_model, state) + + hf_model = transformers.PvtV2ForImageClassification.from_pretrained(hf_id).eval() + results = verify_cls_model_equivalence( + model_a=hf_model, + model_b=keras_model, + input_shape=keras_model.input_shape[1:], + output_specs={"num_classes": keras_model.output_shape[-1]}, + comparison_type="hf_to_keras", + run_performance=False, + atol=1e-2, + rtol=1e-2, + ) + if not results["standard_input"]: + raise ValueError( + "Model equivalence test failed - outputs do not match for standard input" + ) + + out_path = f"{variant}.weights.h5" + keras_model.save_weights(out_path) + print(f" Saved -> {out_path}") + + del keras_model, state, hf_model + keras.backend.clear_session() + gc.collect() diff --git a/zeromodels/models/pvt_v2/pvt_v2_config.py b/zeromodels/models/pvt_v2/pvt_v2_config.py new file mode 100644 index 00000000..79278abc --- /dev/null +++ b/zeromodels/models/pvt_v2/pvt_v2_config.py @@ -0,0 +1,69 @@ +from zeromodels.base import BaseConfig + + +class PvtV2Config(BaseConfig): + r"""Configuration for [`PvtV2Model`] / [`PvtV2ImageClassify`]. + + PVTv2 (Pyramid Vision Transformer v2) is a hierarchical, convolution-augmented + transformer: four stages, each with an OVERLAPPING convolutional patch embedding, + spatial-reduction attention (optionally the linear variant with 7x7 adaptive + pooling), and a convolutional feed-forward network (a 3x3 depthwise conv between the + two Dense layers). It uses no learned position embeddings, so variable input + resolution works out of the box. One `zm_config.json` (declaring the canonical + [`PvtV2ImageClassify`]) sits on each variant's repo; both the backbone and the + classifier load from it. Fields mirror the model constructor and serialize flat. + + Args: + hidden_sizes (`tuple`, *optional*, defaults to `(32, 64, 160, 256)`): + Channel width per stage. + depths (`tuple`, *optional*, defaults to `(2, 2, 2, 2)`): + Number of transformer blocks per stage. + num_attention_heads (`tuple`, *optional*, defaults to `(1, 2, 5, 8)`): + Attention heads per stage. + sr_ratios (`tuple`, *optional*, defaults to `(8, 4, 2, 1)`): + Spatial-reduction ratio of the key/value tokens per stage. + mlp_ratios (`tuple`, *optional*, defaults to `(8, 8, 4, 4)`): + Feed-forward hidden expansion per stage. + linear_attention (`bool`, *optional*, defaults to `False`): + Use the linear-attention variant (7x7 adaptive pooling + 1x1 conv + GELU, + plus a ReLU after the first FFN Dense). + image_size (`int`, *optional*, defaults to 224): + Square input resolution the weights were trained at. + num_classes (`int`, *optional*, defaults to 1000): + Number of classifier output classes (backbone ignores it). + + Examples: + + ```python + >>> from zeromodels.models.pvt_v2 import PvtV2Config, PvtV2ImageClassify + + >>> configuration = PvtV2Config() + >>> model = PvtV2ImageClassify(configuration) + >>> configuration = model.config + ```""" + + model_type = "pvt_v2" + + hidden_sizes: tuple = (32, 64, 160, 256) + depths: tuple = (2, 2, 2, 2) + num_attention_heads: tuple = (1, 2, 5, 8) + sr_ratios: tuple = (8, 4, 2, 1) + mlp_ratios: tuple = (8, 8, 4, 4) + linear_attention: bool = False + image_size: int = 224 + num_classes: int = 1000 + + +# Hosted variants -> arch preset. Weights load by Hub repo id (zm_config.json). +PVT_V2_VARIANTS = { + "pvt_v2_b0": {"model": "pvt_v2_b0", "hf_id": "OpenGVLab/pvt_v2_b0"}, + "pvt_v2_b1": {"model": "pvt_v2_b1", "hf_id": "OpenGVLab/pvt_v2_b1"}, + "pvt_v2_b2": {"model": "pvt_v2_b2", "hf_id": "OpenGVLab/pvt_v2_b2"}, + "pvt_v2_b2_linear": { + "model": "pvt_v2_b2_linear", + "hf_id": "OpenGVLab/pvt_v2_b2_linear", + }, + "pvt_v2_b3": {"model": "pvt_v2_b3", "hf_id": "OpenGVLab/pvt_v2_b3"}, + "pvt_v2_b4": {"model": "pvt_v2_b4", "hf_id": "OpenGVLab/pvt_v2_b4"}, + "pvt_v2_b5": {"model": "pvt_v2_b5", "hf_id": "OpenGVLab/pvt_v2_b5"}, +} diff --git a/zeromodels/models/pvt_v2/pvt_v2_layers.py b/zeromodels/models/pvt_v2/pvt_v2_layers.py new file mode 100644 index 00000000..42f499da --- /dev/null +++ b/zeromodels/models/pvt_v2/pvt_v2_layers.py @@ -0,0 +1,188 @@ +import keras +from keras import layers, ops + +from zeromodels.base.base_attention import fused_attention + + +def to_grid(x, height, width, channels, data_format): + """(B, H*W, C) tokens -> spatial grid in ``data_format`` layout.""" + x = ops.reshape(x, (ops.shape(x)[0], height, width, channels)) + if data_format == "channels_first": + x = ops.transpose(x, (0, 3, 1, 2)) + return x + + +def to_tokens(x, channels, data_format): + """Spatial grid -> (B, H*W, C) tokens.""" + if data_format == "channels_first": + x = ops.transpose(x, (0, 2, 3, 1)) + return ops.reshape(x, (ops.shape(x)[0], -1, channels)) + + +def adaptive_pool_matrix(in_size, out_size): + """Row-averaging matrix ``(out_size, in_size)`` matching torch AdaptiveAvgPool: output + bin ``o`` averages input indices ``[floor(o*in/out), ceil((o+1)*in/out))``.""" + rows = [] + for o in range(out_size): + start = (o * in_size) // out_size + end = -(-(o + 1) * in_size // out_size) # ceil + weight = 1.0 / (end - start) + rows.append([weight if start <= i < end else 0.0 for i in range(in_size)]) + return ops.convert_to_tensor(rows, dtype="float32") + + +@keras.saving.register_keras_serializable(package="zeromodels") +class PvtV2SelfAttention(layers.Layer): + """PVTv2 spatial-reduction attention. + + Query is projected from the full token sequence; keys/values come from a reduced + sequence. Standard SRA (``sr_ratio > 1``) reduces via a strided ``Conv2d`` + LayerNorm; + the linear variant pools the grid to a fixed 7x7 (adaptive average pool), applies a 1x1 + conv + LayerNorm + GELU, so the key/value length is resolution-independent. + """ + + def __init__( + self, + hidden_size, + num_heads, + sr_ratio, + linear_attention=False, + qkv_bias=True, + layer_norm_eps=1e-6, + block_prefix="block", + **kwargs, + ): + super().__init__(**kwargs) + assert hidden_size % num_heads == 0 + self.hidden_size = hidden_size + self.num_heads = num_heads + self.head_dim = hidden_size // num_heads + self.scale = self.head_dim**-0.5 + self.sr_ratio = sr_ratio + self.linear_attention = linear_attention + self.qkv_bias = qkv_bias + self.layer_norm_eps = layer_norm_eps + self.block_prefix = block_prefix + self.data_format = keras.config.image_data_format() + + self.query = layers.Dense( + hidden_size, use_bias=qkv_bias, name=f"{block_prefix}_attn_query" + ) + self.key = layers.Dense( + hidden_size, use_bias=qkv_bias, name=f"{block_prefix}_attn_key" + ) + self.value = layers.Dense( + hidden_size, use_bias=qkv_bias, name=f"{block_prefix}_attn_value" + ) + self.proj = layers.Dense(hidden_size, name=f"{block_prefix}_attn_proj") + + if linear_attention: + self.sr = layers.Conv2D( + hidden_size, + 1, + strides=1, + data_format=self.data_format, + name=f"{block_prefix}_attn_sr", + ) + self.norm = layers.LayerNormalization( + axis=-1, epsilon=layer_norm_eps, name=f"{block_prefix}_attn_norm" + ) + elif sr_ratio > 1: + self.sr = layers.Conv2D( + hidden_size, + sr_ratio, + strides=sr_ratio, + padding="valid", + data_format=self.data_format, + name=f"{block_prefix}_attn_sr", + ) + self.norm = layers.LayerNormalization( + axis=-1, epsilon=layer_norm_eps, name=f"{block_prefix}_attn_norm" + ) + + def split_heads(self, x): + b = ops.shape(x)[0] + x = ops.reshape(x, (b, -1, self.num_heads, self.head_dim)) + return ops.transpose(x, (0, 2, 1, 3)) + + def reduce(self, x, height, width): + grid = to_grid(x, height, width, self.hidden_size, self.data_format) + if self.linear_attention: + grid_cl = ( + ops.transpose(grid, (0, 2, 3, 1)) + if self.data_format == "channels_first" + else grid + ) + mh = adaptive_pool_matrix(height, 7) + mw = adaptive_pool_matrix(width, 7) + grid_cl = ops.einsum("oh,bhwc->bowc", mh, grid_cl) + grid_cl = ops.einsum("pw,bowc->bopc", mw, grid_cl) + grid = ( + ops.transpose(grid_cl, (0, 3, 1, 2)) + if self.data_format == "channels_first" + else grid_cl + ) + x = to_tokens(self.sr(grid), self.hidden_size, self.data_format) + x = ops.gelu(self.norm(x)) + else: + x = to_tokens(self.sr(grid), self.hidden_size, self.data_format) + x = self.norm(x) + return x + + def call(self, x, height, width, training=None): + q = self.split_heads(self.query(x)) + if self.linear_attention or self.sr_ratio > 1: + kv_in = self.reduce(x, height, width) + else: + kv_in = x + k = self.split_heads(self.key(kv_in)) + v = self.split_heads(self.value(kv_in)) + + out = fused_attention(q, k, v, self.scale, training=training) + out = ops.transpose(out, (0, 2, 1, 3)) + out = ops.reshape(out, (ops.shape(x)[0], ops.shape(x)[1], self.hidden_size)) + return self.proj(out) + + def compute_output_shape(self, input_shape): + return (input_shape[0], input_shape[1], self.hidden_size) + + def get_config(self): + config = super().get_config() + config.update( + { + "hidden_size": self.hidden_size, + "num_heads": self.num_heads, + "sr_ratio": self.sr_ratio, + "linear_attention": self.linear_attention, + "qkv_bias": self.qkv_bias, + "layer_norm_eps": self.layer_norm_eps, + "block_prefix": self.block_prefix, + } + ) + return config + + +@keras.saving.register_keras_serializable(package="zeromodels") +class PvtDropPath(layers.Layer): + """Stochastic depth: drops the residual branch per-sample during training only.""" + + def __init__(self, drop_prob, seed=None, **kwargs): + super().__init__(**kwargs) + self.drop_prob = drop_prob + self.seed = seed + self.seed_generator = keras.random.SeedGenerator(seed) + + def call(self, x, training=None): + if training and self.drop_prob > 0: + keep = 1 - self.drop_prob + shape = (ops.shape(x)[0],) + (1,) * (len(x.shape) - 1) + mask = ops.floor( + keep + keras.random.uniform(shape, 0, 1, seed=self.seed_generator) + ) + return (x / keep) * mask + return x + + def get_config(self): + config = super().get_config() + config.update({"drop_prob": self.drop_prob, "seed": self.seed}) + return config diff --git a/zeromodels/models/pvt_v2/pvt_v2_model.py b/zeromodels/models/pvt_v2/pvt_v2_model.py new file mode 100644 index 00000000..4e097403 --- /dev/null +++ b/zeromodels/models/pvt_v2/pvt_v2_model.py @@ -0,0 +1,423 @@ +import keras +from keras import layers + +from zeromodels.base import BaseModel +from zeromodels.conversion import copy_weights_by_path_suffix +from zeromodels.utils import standardize_input_shape +from zeromodels.utils.image_util import normalize_image_for_classify_models + +from .pvt_v2_config import PvtV2Config +from .pvt_v2_layers import PvtDropPath, PvtV2SelfAttention + +PVT_V2_HUB_SIBLINGS = frozenset({"PvtV2Model", "PvtV2ImageClassify"}) + +PATCH_SIZES = (7, 3, 3, 3) +STRIDES = (4, 2, 2, 2) + + +def tokens_to_grid(x, H, W, channels, data_format): + """(B, H*W, C) tokens -> spatial grid in ``data_format`` layout (functional).""" + x = layers.Reshape((H, W, channels))(x) + if data_format == "channels_first": + x = keras.ops.transpose(x, (0, 3, 1, 2)) + return x + + +def grid_to_tokens(x, channels, data_format): + """Spatial grid -> (B, H*W, C) tokens (functional).""" + if data_format == "channels_first": + x = keras.ops.transpose(x, (0, 2, 3, 1)) + return layers.Reshape((-1, channels))(x) + + +def overlap_patch_embed(x, out_channels, patch_size, stride, data_format, stage_idx): + """Overlapping patch embed: symmetric ZeroPad -> Conv(valid) -> tokens -> LayerNorm. + Returns ``(tokens, H, W)``.""" + x = layers.ZeroPadding2D(padding=patch_size // 2, data_format=data_format)(x) + x = layers.Conv2D( + out_channels, + patch_size, + strides=stride, + padding="valid", + data_format=data_format, + name=f"layers_{stage_idx}_patch_embed_proj", + )(x) + if data_format == "channels_first": + H, W = int(x.shape[2]), int(x.shape[3]) + else: + H, W = int(x.shape[1]), int(x.shape[2]) + x = grid_to_tokens(x, out_channels, data_format) + x = layers.LayerNormalization( + axis=-1, epsilon=1e-6, name=f"layers_{stage_idx}_patch_embed_layernorm" + )(x) + return x, H, W + + +def conv_mlp(x, H, W, channels, mid_channels, linear, data_format, name_prefix): + """PVTv2 conv-FFN: Dense -> (ReLU if linear) -> DWConv -> GELU -> Dense.""" + x = layers.Dense(mid_channels, name=f"{name_prefix}_dense1")(x) + if linear: + x = layers.Activation("relu")(x) + grid = tokens_to_grid(x, H, W, mid_channels, data_format) + grid = layers.DepthwiseConv2D( + 3, + strides=1, + padding="same", + data_format=data_format, + name=f"{name_prefix}_dwconv", + )(grid) + x = grid_to_tokens(grid, mid_channels, data_format) + x = layers.Activation("gelu")(x) + x = layers.Dense(channels, name=f"{name_prefix}_dense2")(x) + return x + + +def pvt_v2_block( + x, + H, + W, + dim, + num_heads, + sr_ratio, + mlp_ratio, + linear, + drop_prob, + data_format, + stage_idx, + block_idx, +): + prefix = f"layers_{stage_idx}_blocks_{block_idx}" + drop_path = PvtDropPath(drop_prob) + + norm1 = layers.LayerNormalization( + axis=-1, epsilon=1e-6, name=f"{prefix}_layernorm_1" + )(x) + attn = PvtV2SelfAttention( + dim, + num_heads, + sr_ratio, + linear_attention=linear, + qkv_bias=True, + block_prefix=prefix, + )(norm1, height=H, width=W) + x = layers.Add()([x, drop_path(attn)]) + + norm2 = layers.LayerNormalization( + axis=-1, epsilon=1e-6, name=f"{prefix}_layernorm_2" + )(x) + mlp = conv_mlp( + norm2, + H, + W, + channels=dim, + mid_channels=int(dim * mlp_ratio), + linear=linear, + data_format=data_format, + name_prefix=f"{prefix}_mlp", + ) + return layers.Add()([x, drop_path(mlp)]) + + +def pvt_v2_backbone_feature( + inputs, + *, + hidden_sizes, + depths, + num_attention_heads, + sr_ratios, + mlp_ratios, + linear_attention, + drop_path_rate, + data_format, + return_stages=False, +): + total = sum(depths) + dpr = [drop_path_rate * i / max(total - 1, 1) for i in range(total)] + x = inputs + features = [] + cur = 0 + for i in range(4): + x, H, W = overlap_patch_embed( + x, hidden_sizes[i], PATCH_SIZES[i], STRIDES[i], data_format, i + ) + for j in range(depths[i]): + x = pvt_v2_block( + x, + H, + W, + hidden_sizes[i], + num_attention_heads[i], + sr_ratios[i], + mlp_ratios[i], + linear_attention, + dpr[cur], + data_format, + i, + j, + ) + cur += 1 + x = layers.LayerNormalization( + axis=-1, epsilon=1e-6, name=f"layers_{i}_layernorm" + )(x) + x = tokens_to_grid(x, H, W, hidden_sizes[i], data_format) + features.append(x) + return features if return_stages else features[-1] + + +@keras.saving.register_keras_serializable(package="zeromodels") +class PvtV2Model(BaseModel): + """Pyramid Vision Transformer v2 (PVTv2) backbone. + + A hierarchical, convolution-augmented transformer: four stages, each an overlapping + convolutional patch embedding, spatial-reduction (or linear) attention, and a + convolutional feed-forward network (a 3x3 depthwise conv between the two Dense + layers). It uses no learned position embeddings, so variable input resolution works + out of the box. Output is the final stage's spatial feature map, or, with + ``as_backbone=True``, the four per-stage feature maps. + + References: + - [PVTv2: Improved Baselines with Pyramid Vision Transformer](https://arxiv.org/abs/2106.13797) + + Args: + See :class:`PvtV2Config`. ``as_backbone`` returns the 4-stage pyramid; + ``include_normalization`` bakes ImageNet normalization into the graph; + ``image_size`` sets the input the model is built for. Defaults describe PVTv2-B0. + """ + + BASE_WEIGHT_CONFIG = None + config_class = PvtV2Config + HUB_REPO_SIBLINGS = PVT_V2_HUB_SIBLINGS + HF_MODEL_TYPE = "pvt_v2" + + @classmethod + def from_hub_repo(cls, repo_id, load_weights=True, skip_mismatch=False, **kwargs): + model = cls.build_from_hub_repo(repo_id, **kwargs) + if load_weights: + src = PvtV2ImageClassify.from_weights(repo_id, skip_mismatch=skip_mismatch) + copy_weights_by_path_suffix(src, model) + del src + return model + + @classmethod + def config_from_hf(cls, hf_config): + return { + "hidden_sizes": hf_config["hidden_sizes"], + "depths": hf_config["depths"], + "num_attention_heads": hf_config["num_attention_heads"], + "sr_ratios": hf_config["sr_ratios"], + "mlp_ratios": hf_config["mlp_ratios"], + "linear_attention": hf_config.get("linear_attention", False), + } + + @classmethod + def transfer_from_hf(cls, keras_model, state_dict): + from .convert_pvt_v2_hf_to_keras import transfer_pvt_v2_weights + + transfer_pvt_v2_weights(keras_model, state_dict) + + def __init__( + self, + as_backbone=False, + hidden_sizes=(32, 64, 160, 256), + depths=(2, 2, 2, 2), + num_attention_heads=(1, 2, 5, 8), + sr_ratios=(8, 4, 2, 1), + mlp_ratios=(8, 8, 4, 4), + linear_attention=False, + drop_path_rate=0.0, + image_size=224, + include_normalization=True, + normalization_mode="imagenet", + input_tensor=None, + name="PvtV2Model", + **kwargs, + ): + for k in ("num_classes", "classifier_activation", "hf_id"): + kwargs.pop(k, None) + + data_format = keras.config.image_data_format() + image_size = standardize_input_shape(image_size, data_format) + + if input_tensor is None: + img_input = layers.Input(shape=image_size) + elif not keras.utils.is_keras_tensor(input_tensor): + img_input = layers.Input(tensor=input_tensor, shape=image_size) + else: + img_input = input_tensor + + x = ( + normalize_image_for_classify_models(img_input, normalization_mode) + if include_normalization + else img_input + ) + features = pvt_v2_backbone_feature( + x, + hidden_sizes=hidden_sizes, + depths=depths, + num_attention_heads=num_attention_heads, + sr_ratios=sr_ratios, + mlp_ratios=mlp_ratios, + linear_attention=linear_attention, + drop_path_rate=drop_path_rate, + data_format=data_format, + return_stages=as_backbone, + ) + super().__init__(inputs=img_input, outputs=features, name=name, **kwargs) + + self.as_backbone = as_backbone + self.hidden_sizes = list(hidden_sizes) + self.depths = list(depths) + self.num_attention_heads = list(num_attention_heads) + self.sr_ratios = list(sr_ratios) + self.mlp_ratios = list(mlp_ratios) + self.linear_attention = linear_attention + self.drop_path_rate = drop_path_rate + self.image_size = image_size + self.include_normalization = include_normalization + self.normalization_mode = normalization_mode + self.input_tensor = input_tensor + + def get_config(self): + config = super().get_config() + config.update( + { + "as_backbone": self.as_backbone, + "hidden_sizes": self.hidden_sizes, + "depths": self.depths, + "num_attention_heads": self.num_attention_heads, + "sr_ratios": self.sr_ratios, + "mlp_ratios": self.mlp_ratios, + "linear_attention": self.linear_attention, + "drop_path_rate": self.drop_path_rate, + "image_size": self.image_size, + "include_normalization": self.include_normalization, + "normalization_mode": self.normalization_mode, + "input_tensor": self.input_tensor, + "name": self.name, + } + ) + return config + + @classmethod + def from_config(cls, config): + return cls(**config) + + +@keras.saving.register_keras_serializable(package="zeromodels") +class PvtV2ImageClassify(BaseModel): + """PVTv2 image classifier: :class:`PvtV2Model` backbone + global-average-pool + + a single Dense head over the last stage's feature map. + + References: + - [PVTv2: Improved Baselines with Pyramid Vision Transformer](https://arxiv.org/abs/2106.13797) + + Args: + See :class:`PvtV2Config`. ``num_classes`` / ``classifier_activation`` are + head-specific; all other args forward to :class:`PvtV2Model`. + """ + + BASE_WEIGHT_CONFIG = None + config_class = PvtV2Config + HUB_REPO_SIBLINGS = PVT_V2_HUB_SIBLINGS + HF_MODEL_TYPE = "pvt_v2" + + @classmethod + def config_from_hf(cls, hf_config): + return { + "hidden_sizes": hf_config["hidden_sizes"], + "depths": hf_config["depths"], + "num_attention_heads": hf_config["num_attention_heads"], + "sr_ratios": hf_config["sr_ratios"], + "mlp_ratios": hf_config["mlp_ratios"], + "linear_attention": hf_config.get("linear_attention", False), + "num_classes": hf_config.get("num_labels", 1000), + } + + @classmethod + def transfer_from_hf(cls, keras_model, state_dict): + from .convert_pvt_v2_hf_to_keras import transfer_pvt_v2_weights + + transfer_pvt_v2_weights(keras_model, state_dict) + + def __init__( + self, + hidden_sizes=(32, 64, 160, 256), + depths=(2, 2, 2, 2), + num_attention_heads=(1, 2, 5, 8), + sr_ratios=(8, 4, 2, 1), + mlp_ratios=(8, 8, 4, 4), + linear_attention=False, + drop_path_rate=0.0, + image_size=224, + include_normalization=True, + normalization_mode="imagenet", + input_tensor=None, + num_classes=1000, + classifier_activation="linear", + name="PvtV2ImageClassify", + **kwargs, + ): + kwargs.pop("hf_id", None) + data_format = keras.config.image_data_format() + + backbone = PvtV2Model( + hidden_sizes=hidden_sizes, + depths=depths, + num_attention_heads=num_attention_heads, + sr_ratios=sr_ratios, + mlp_ratios=mlp_ratios, + linear_attention=linear_attention, + drop_path_rate=drop_path_rate, + image_size=image_size, + include_normalization=include_normalization, + normalization_mode=normalization_mode, + input_tensor=input_tensor, + name=f"{name}_backbone", + ) + x = layers.GlobalAveragePooling2D(data_format=data_format, name="avg_pool")( + backbone.output + ) + out = layers.Dense( + num_classes, activation=classifier_activation, name="predictions" + )(x) + super().__init__(inputs=backbone.input, outputs=out, name=name, **kwargs) + + self.hidden_sizes = list(hidden_sizes) + self.depths = list(depths) + self.num_attention_heads = list(num_attention_heads) + self.sr_ratios = list(sr_ratios) + self.mlp_ratios = list(mlp_ratios) + self.linear_attention = linear_attention + self.drop_path_rate = drop_path_rate + self.image_size = backbone.image_size + self.include_normalization = include_normalization + self.normalization_mode = normalization_mode + self.input_tensor = input_tensor + self.num_classes = num_classes + self.classifier_activation = classifier_activation + + def get_config(self): + config = super().get_config() + config.update( + { + "hidden_sizes": self.hidden_sizes, + "depths": self.depths, + "num_attention_heads": self.num_attention_heads, + "sr_ratios": self.sr_ratios, + "mlp_ratios": self.mlp_ratios, + "linear_attention": self.linear_attention, + "drop_path_rate": self.drop_path_rate, + "image_size": self.image_size, + "include_normalization": self.include_normalization, + "normalization_mode": self.normalization_mode, + "input_tensor": self.input_tensor, + "num_classes": self.num_classes, + "classifier_activation": self.classifier_activation, + "name": self.name, + } + ) + return config + + @classmethod + def from_config(cls, config): + return cls(**config)