kmodels 0.2.0__tar.gz → 0.2.2__tar.gz

{kmodels-0.2.0 → kmodels-0.2.2}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: kmodels
-Version: 0.2.0
+Version: 0.2.2
 Summary: Pretrained keras 3 vision models
 Author-email: Gitesh Chawda <gitesh.ch.0912@gmail.com>
 License: Apache License 2.0
@@ -25,6 +25,10 @@ Requires-Python: >=3.11
 Description-Content-Type: text/markdown
 License-File: LICENSE
 Requires-Dist: keras
+Provides-Extra: test
+Requires-Dist: pytest; extra == "test"
+Requires-Dist: pytest-cov; extra == "test"
+Requires-Dist: requests; extra == "test"
 Dynamic: license-file
 
 # Keras Models 🚀
@@ -35,7 +39,7 @@ Dynamic: license-file
 
 ## 📖 Introduction
 
-Keras Models (kmodels) is a collection of models with pretrained weights, built entirely with Keras 3. It supports a range of tasks, including classification, object detection (DETR), segmentation (SAM, SegFormer, DeepLabV3, EoMT), vision-language modeling (CLIP, SigLIP, SigLIP2), and more. kmodels includes custom layers and backbone support, providing flexibility and efficiency across various applications. For backbones, there are various weight variants like `in1k`, `in21k`, `fb_dist_in1k`, `ms_in22k`, `fb_in22k_ft_in1k`, `ns_jft_in1k`, `aa_in1k`, `cvnets_in1k`, `augreg_in21k_ft_in1k`, `augreg_in21k`, and many more.
+Keras Models (kmodels) is a collection of models with pretrained weights, built entirely with Keras 3. It supports a range of tasks, including classification, object detection (DETR, RT-DETR, RF-DETR), segmentation (SAM, SAM2, SegFormer, DeepLabV3, EoMT), vision-language modeling (CLIP, SigLIP, SigLIP2), and more. kmodels includes custom layers and backbone support, providing flexibility and efficiency across various applications. For backbones, there are various weight variants like `in1k`, `in21k`, `fb_dist_in1k`, `ms_in22k`, `fb_in22k_ft_in1k`, `ns_jft_in1k`, `aa_in1k`, `cvnets_in1k`, `augreg_in21k_ft_in1k`, `augreg_in21k`, and many more.
 
 ## ⚡ Installation
 
@@ -62,6 +66,7 @@ pip install -U git+https://github.com/IMvision12/keras-models
 | Model | Description |
 |-------|-------------|
 | [SAM](docs/sam.md) | Segment Anything Model — promptable segmentation with points, boxes, or masks (ViT-B/L/H) |
+| [SAM2](docs/sam2.md) | Segment Anything Model 2 — next generation of promptable visual segmentation (Hiera Tiny/Small/Base+/Large) |
 | [SegFormer](docs/segformer.md) | Transformer-based semantic segmentation with MLP decoder, Cityscapes & ADE20K weights |
 | [DeepLabV3](docs/deeplabv3.md) | Atrous convolution-based semantic segmentation |
 | [EoMT](docs/eomt.md) | Encoder-only Mask Transformer for panoptic segmentation |
@@ -71,6 +76,7 @@ pip install -U git+https://github.com/IMvision12/keras-models
 | Model | Description |
 |-------|-------------|
 | [DETR](docs/detr.md) | End-to-end object detection with Transformers (ResNet-50/101 backbones) |
+| [RT-DETR](docs/rt_detr.md) | Real-time DETR with ResNet-vd backbone and hybrid encoder (ResNet-18/34/50/101 variants) |
 | [RF-DETR](docs/rf_detr.md) | Real-time detection transformer (Nano, Small, Medium, Base, Large variants) |
 
 **Vision-Language Models**
@@ -107,6 +113,7 @@ pip install -U git+https://github.com/IMvision12/keras-models
     | MobileNetV3 | [Searching for MobileNetV3](https://arxiv.org/abs/1905.02244) | `keras` |
     | MobileViT | [MobileViT: Light-weight, General-purpose, and Mobile-friendly Vision Transformer](https://arxiv.org/abs/2110.02178) | `timm` |
     | MobileViTV2 | [Separable Self-attention for Mobile Vision Transformers](https://arxiv.org/abs/2206.02680) | `timm` |
+    | 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` |
     | Res2Net | [Res2Net: A New Multi-scale Backbone Architecture](https://arxiv.org/abs/1904.01169) | `timm` |
@@ -127,6 +134,7 @@ pip install -U git+https://github.com/IMvision12/keras-models
     | 🏷️ Model Name | 📜 Reference Paper | 📦 Source of Weights |
     |---------------|-------------------|---------------------|
     | DETR | [End-to-End Object Detection with Transformers](https://arxiv.org/abs/2005.12872) | `transformers`|
+    | RT-DETR | [DETRs Beat YOLOs on Real-time Object Detection](https://arxiv.org/abs/2304.08069) | `transformers` |
     | RF-DETR | [RF-DETR: Real-Time Detection Transformer](https://arxiv.org/abs/2502.18860) | `rfdetr` |
 
 <br>
@@ -138,6 +146,7 @@ pip install -U git+https://github.com/IMvision12/keras-models
     | DeepLabV3 | [Rethinking Atrous Convolution for Semantic Image Segmentation](https://arxiv.org/abs/1706.05587) | `torchvision` |
     | EoMT | [Encoder-only Mask Transformer for Panoptic Segmentation](https://arxiv.org/abs/2504.07957) | `transformers` |
     | SAM | [Segment Anything](https://arxiv.org/abs/2304.02643) | `transformers` |
+    | SAM2 | [SAM 2: Segment Anything in Images and Videos](https://arxiv.org/abs/2408.00714) | `transformers` |
     | SegFormer | [SegFormer: Simple and Efficient Design for Semantic Segmentation with Transformers](https://arxiv.org/abs/2105.15203) | `transformers`|
 
 <br>

{kmodels-0.2.0 → kmodels-0.2.2}/README.md

@@ -6,7 +6,7 @@
 
 ## 📖 Introduction
 
-Keras Models (kmodels) is a collection of models with pretrained weights, built entirely with Keras 3. It supports a range of tasks, including classification, object detection (DETR), segmentation (SAM, SegFormer, DeepLabV3, EoMT), vision-language modeling (CLIP, SigLIP, SigLIP2), and more. kmodels includes custom layers and backbone support, providing flexibility and efficiency across various applications. For backbones, there are various weight variants like `in1k`, `in21k`, `fb_dist_in1k`, `ms_in22k`, `fb_in22k_ft_in1k`, `ns_jft_in1k`, `aa_in1k`, `cvnets_in1k`, `augreg_in21k_ft_in1k`, `augreg_in21k`, and many more.
+Keras Models (kmodels) is a collection of models with pretrained weights, built entirely with Keras 3. It supports a range of tasks, including classification, object detection (DETR, RT-DETR, RF-DETR), segmentation (SAM, SAM2, SegFormer, DeepLabV3, EoMT), vision-language modeling (CLIP, SigLIP, SigLIP2), and more. kmodels includes custom layers and backbone support, providing flexibility and efficiency across various applications. For backbones, there are various weight variants like `in1k`, `in21k`, `fb_dist_in1k`, `ms_in22k`, `fb_in22k_ft_in1k`, `ns_jft_in1k`, `aa_in1k`, `cvnets_in1k`, `augreg_in21k_ft_in1k`, `augreg_in21k`, and many more.
 
 ## ⚡ Installation
 
@@ -33,6 +33,7 @@ pip install -U git+https://github.com/IMvision12/keras-models
 | Model | Description |
 |-------|-------------|
 | [SAM](docs/sam.md) | Segment Anything Model — promptable segmentation with points, boxes, or masks (ViT-B/L/H) |
+| [SAM2](docs/sam2.md) | Segment Anything Model 2 — next generation of promptable visual segmentation (Hiera Tiny/Small/Base+/Large) |
 | [SegFormer](docs/segformer.md) | Transformer-based semantic segmentation with MLP decoder, Cityscapes & ADE20K weights |
 | [DeepLabV3](docs/deeplabv3.md) | Atrous convolution-based semantic segmentation |
 | [EoMT](docs/eomt.md) | Encoder-only Mask Transformer for panoptic segmentation |
@@ -42,6 +43,7 @@ pip install -U git+https://github.com/IMvision12/keras-models
 | Model | Description |
 |-------|-------------|
 | [DETR](docs/detr.md) | End-to-end object detection with Transformers (ResNet-50/101 backbones) |
+| [RT-DETR](docs/rt_detr.md) | Real-time DETR with ResNet-vd backbone and hybrid encoder (ResNet-18/34/50/101 variants) |
 | [RF-DETR](docs/rf_detr.md) | Real-time detection transformer (Nano, Small, Medium, Base, Large variants) |
 
 **Vision-Language Models**
@@ -78,6 +80,7 @@ pip install -U git+https://github.com/IMvision12/keras-models
     | MobileNetV3 | [Searching for MobileNetV3](https://arxiv.org/abs/1905.02244) | `keras` |
     | MobileViT | [MobileViT: Light-weight, General-purpose, and Mobile-friendly Vision Transformer](https://arxiv.org/abs/2110.02178) | `timm` |
     | MobileViTV2 | [Separable Self-attention for Mobile Vision Transformers](https://arxiv.org/abs/2206.02680) | `timm` |
+    | 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` |
     | Res2Net | [Res2Net: A New Multi-scale Backbone Architecture](https://arxiv.org/abs/1904.01169) | `timm` |
@@ -98,6 +101,7 @@ pip install -U git+https://github.com/IMvision12/keras-models
     | 🏷️ Model Name | 📜 Reference Paper | 📦 Source of Weights |
     |---------------|-------------------|---------------------|
     | DETR | [End-to-End Object Detection with Transformers](https://arxiv.org/abs/2005.12872) | `transformers`|
+    | RT-DETR | [DETRs Beat YOLOs on Real-time Object Detection](https://arxiv.org/abs/2304.08069) | `transformers` |
     | RF-DETR | [RF-DETR: Real-Time Detection Transformer](https://arxiv.org/abs/2502.18860) | `rfdetr` |
 
 <br>
@@ -109,6 +113,7 @@ pip install -U git+https://github.com/IMvision12/keras-models
     | DeepLabV3 | [Rethinking Atrous Convolution for Semantic Image Segmentation](https://arxiv.org/abs/1706.05587) | `torchvision` |
     | EoMT | [Encoder-only Mask Transformer for Panoptic Segmentation](https://arxiv.org/abs/2504.07957) | `transformers` |
     | SAM | [Segment Anything](https://arxiv.org/abs/2304.02643) | `transformers` |
+    | SAM2 | [SAM 2: Segment Anything in Images and Videos](https://arxiv.org/abs/2408.00714) | `transformers` |
     | SegFormer | [SegFormer: Simple and Efficient Design for Semantic Segmentation with Transformers](https://arxiv.org/abs/2105.15203) | `transformers`|
 
 <br>

{kmodels-0.2.0 → kmodels-0.2.2}/kmodels/__init__.py

@@ -2,4 +2,4 @@ from kmodels import layers, models, utils
 from kmodels.model_registry import list_models, register_model
 from kmodels.version import version
 
-__version__ = "0.2.0"
+__version__ = "0.2.2"

kmodels-0.2.2/kmodels/_test_runner.py

@@ -0,0 +1,171 @@
+"""Cross-platform test runner for keras-models.
+
+Usage:
+    kmodels-test <command>
+
+Commands:
+    all                  Full test suite (torch, excludes slow/link tests)
+    backend-torch        Backend tests on torch
+    backend-jax          Backend tests on jax
+    backend-tf           Backend tests on tensorflow
+    sas-torch            Serialization + saving on torch
+    sas-tf               Serialization + saving on tensorflow
+    sas-jax              Serialization + saving on jax
+    df-torch             Data format tests on torch
+    df-tf                Data format tests on tensorflow (GPU auto-skip)
+    df-jax               Data format tests on jax
+    gpu                  GPU-marked tests only
+    gpu-all              Full test suite on GPU (torch + tf)
+    help                 Show this message
+"""
+
+import os
+import subprocess
+import sys
+
+PYTEST = [sys.executable, "-m", "pytest"]
+
+
+def _run(backend, *pytest_args):
+    """Run pytest with the given backend and arguments."""
+    env = os.environ.copy()
+    if backend:
+        env["KERAS_BACKEND"] = backend
+    cmd = PYTEST + list(pytest_args)
+    print(f"\n{'=' * 60}")
+    print(f"  KERAS_BACKEND={backend or '(default)'}  {' '.join(pytest_args)}")
+    print(f"{'=' * 60}\n")
+    result = subprocess.run(cmd, env=env)
+    return result.returncode
+
+
+COMMANDS = {}
+
+
+def command(name):
+    def decorator(fn):
+        COMMANDS[name] = fn
+        return fn
+
+    return decorator
+
+
+@command("all")
+def test_all():
+    return _run(
+        "torch",
+        "tests/",
+        "-v",
+        "--durations=20",
+        "-m",
+        "not slow and not gpu",
+    )
+
+
+@command("backend-torch")
+def test_backend_torch():
+    return _run("torch", "tests/integration/test_backend_compatibility.py", "-v")
+
+
+@command("backend-jax")
+def test_backend_jax():
+    return _run("jax", "tests/integration/test_backend_compatibility.py", "-v")
+
+
+@command("backend-tf")
+def test_backend_tf():
+    return _run("tensorflow", "tests/integration/test_backend_compatibility.py", "-v")
+
+
+SAS_FILES = [
+    "tests/integration/test_serialization.py",
+    "tests/integration/test_model_saving.py",
+]
+
+
+@command("sas-torch")
+def test_sas_torch():
+    return _run("torch", *SAS_FILES, "-v")
+
+
+@command("sas-tf")
+def test_sas_tf():
+    return _run("tensorflow", *SAS_FILES, "-v")
+
+
+@command("sas-jax")
+def test_sas_jax():
+    return _run("jax", *SAS_FILES, "-v")
+
+
+DF_FILE = "tests/integration/test_data_formats.py"
+
+
+@command("df-torch")
+def test_df_torch():
+    return _run("torch", DF_FILE, "-v")
+
+
+@command("df-tf")
+def test_df_tf():
+    return _run("tensorflow", DF_FILE, "-v")
+
+
+@command("df-jax")
+def test_df_jax():
+    return _run("jax", DF_FILE, "-v")
+
+
+@command("gpu")
+def test_gpu():
+    rc1 = _run("torch", "tests/", "-v", "-m", "gpu")
+    rc2 = _run(
+        "tensorflow",
+        "tests/integration/test_data_formats.py",
+        "-v",
+        "-k",
+        "channels_first",
+    )
+    return rc1 or rc2
+
+
+@command("gpu-all")
+def test_gpu_all():
+    rc1 = _run(
+        "torch",
+        "tests/",
+        "-v",
+        "--durations=20",
+        "-m",
+        "not slow and not link_validation",
+    )
+    rc2 = _run(
+        "tensorflow",
+        "tests/",
+        "-v",
+        "--durations=20",
+        "-m",
+        "not slow and not link_validation",
+    )
+    return rc1 or rc2
+
+
+@command("help")
+def show_help():
+    print(__doc__)
+    return 0
+
+
+def main():
+    if len(sys.argv) < 2 or sys.argv[1] not in COMMANDS:
+        show_help()
+        if len(sys.argv) >= 2 and sys.argv[1] not in COMMANDS:
+            print(f"\nError: Unknown command '{sys.argv[1]}'")
+            return 1
+        return 0
+
+    return COMMANDS[sys.argv[1]]()
+
+
+if __name__ == "__main__":
+    sys.exit(main())

{kmodels-0.2.0 → kmodels-0.2.2}/kmodels/models/__init__.py

@@ -24,6 +24,7 @@ from kmodels.models import (
     mobilenetv3,
     mobilevit,
     mobilevitv2,
+    nextvit,
     pit,
     poolformer,
     res2net,
@@ -32,7 +33,9 @@ from kmodels.models import (
     resnetv2,
     resnext,
     rf_detr,
+    rt_detr,
     sam,
+    sam2,
     segformer,
     senet,
     siglip,

{kmodels-0.2.0 → kmodels-0.2.2}/kmodels/models/cait/cait_layers.py

@@ -81,7 +81,7 @@ class ClassDistToken(layers.Layer):
                     [cls_broadcasted, dist_broadcasted, inputs], axis=1
                 )
             else:
-                return ops.concatenate([cls_broadcasted, inputs], axis=1)
+                return cls_broadcasted
 
     def get_config(self):
         config = super().get_config()

{kmodels-0.2.0 → kmodels-0.2.2}/kmodels/models/cait/cait_model.py

@@ -313,8 +313,12 @@ class CaiT(keras.Model):
             name="stem_conv",
         )(x)
 
-        grid_h = input_shape[0] // patch_size
-        grid_w = input_shape[1] // patch_size
+        if data_format == "channels_first":
+            grid_h = input_shape[1] // patch_size
+            grid_w = input_shape[2] // patch_size
+        else:
+            grid_h = input_shape[0] // patch_size
+            grid_w = input_shape[1] // patch_size
 
         x = layers.Reshape((-1, embed_dim))(x)
 
@@ -351,7 +355,8 @@ class CaiT(keras.Model):
             if i == depth_token_only - 1:
                 features.append(cls_token)
 
-        x = layers.LayerNormalization(epsilon=1e-6, name="final_layernorm")(cls_token)
+        x = layers.Concatenate(axis=1, name="cat_cls_patch")([cls_token, x])
+        x = layers.LayerNormalization(epsilon=1e-6, name="final_layernorm")(x)
 
         if include_top:
             x = layers.Dense(

kmodels-0.2.2/kmodels/models/cait/convert_cait_torch_to_keras.py

@@ -0,0 +1,205 @@
+import re
+from typing import Dict, List, Union
+
+import keras
+import timm
+import torch
+from tqdm import tqdm
+
+from kmodels.models import cait
+from kmodels.utils.custom_exception import WeightMappingError, WeightShapeMismatchError
+from kmodels.utils.model_equivalence_tester import verify_cls_model_equivalence
+from kmodels.utils.weight_split_torch_and_keras import split_model_weights
+from kmodels.utils.weight_transfer_torch_to_keras import (
+    compare_keras_torch_names,
+    transfer_attention_weights,
+    transfer_weights,
+)
+
+weight_name_mapping = {
+    "_": ".",
+    "stem.conv": "patch_embed.proj",
+    "cls.token.cls.token": "cls_token",
+    "pos.embed.pos.embed": "pos_embed",
+    "layernorm.": "norm",
+    "dense.1": "fc1",
+    "dense.2": "fc2",
+    "blocks.token.only": "blocks_token_only",
+    "kernel": "weight",
+    "gamma": "weight",
+    "beta": "bias",
+    "moving_mean": "running_mean",
+    "moving_variance": "running_var",
+    "final.norm": "norm.",
+    "predictions": "head",
+}
+
+attn_weight_replacement: Dict[str, str] = {
+    "proj.l": "proj_l",
+    "proj.w": "proj_w",
+    "blocks.token.only": "blocks_token_only",
+}
+
+model_configs: List[Dict[str, Union[str, type]]] = [
+    {
+        "keras_cls": cait.CaiTXXS24,
+        "torch_name": "cait_xxs24_224.fb_dist_in1k",
+        "input_shape": [224, 224, 3],
+        "num_classes": 1000,
+    },
+    {
+        "keras_cls": cait.CaiTXXS24,
+        "torch_name": "cait_xxs24_384.fb_dist_in1k",
+        "input_shape": [384, 384, 3],
+        "num_classes": 1000,
+    },
+    {
+        "keras_cls": cait.CaiTXXS36,
+        "torch_name": "cait_xxs36_224.fb_dist_in1k",
+        "input_shape": [224, 224, 3],
+        "num_classes": 1000,
+    },
+    {
+        "keras_cls": cait.CaiTXXS36,
+        "torch_name": "cait_xxs36_384.fb_dist_in1k",
+        "input_shape": [384, 384, 3],
+        "num_classes": 1000,
+    },
+    {
+        "keras_cls": cait.CaiTXS24,
+        "torch_name": "cait_xs24_384.fb_dist_in1k",
+        "input_shape": [384, 384, 3],
+        "num_classes": 1000,
+    },
+    {
+        "keras_cls": cait.CaiTS24,
+        "torch_name": "cait_s24_224.fb_dist_in1k",
+        "input_shape": [224, 224, 3],
+        "num_classes": 1000,
+    },
+    {
+        "keras_cls": cait.CaiTS24,
+        "torch_name": "cait_s24_384.fb_dist_in1k",
+        "input_shape": [384, 384, 3],
+        "num_classes": 1000,
+    },
+    {
+        "keras_cls": cait.CaiTS36,
+        "torch_name": "cait_s36_384.fb_dist_in1k",
+        "input_shape": [384, 384, 3],
+        "num_classes": 1000,
+    },
+    {
+        "keras_cls": cait.CaiTM36,
+        "torch_name": "cait_m36_384.fb_dist_in1k",
+        "input_shape": [384, 384, 3],
+        "num_classes": 1000,
+    },
+    {
+        "keras_cls": cait.CaiTM48,
+        "torch_name": "cait_m48_448.fb_dist_in1k",
+        "input_shape": [448, 448, 3],
+        "num_classes": 1000,
+    },
+]
+
+for model_config in model_configs:
+    torch_model_name: str = model_config["torch_name"]
+    print(f"\n{'=' * 60}")
+    print(f"Converting {torch_model_name}...")
+    print(f"{'=' * 60}")
+
+    keras_model: keras.Model = model_config["keras_cls"](
+        include_top=True,
+        input_shape=model_config["input_shape"],
+        classifier_activation="linear",
+        num_classes=model_config["num_classes"],
+        include_normalization=False,
+        weights=None,
+    )
+
+    torch_model: torch.nn.Module = timm.create_model(
+        torch_model_name, pretrained=True
+    ).eval()
+
+    trainable_torch_weights, non_trainable_torch_weights, _ = split_model_weights(
+        torch_model
+    )
+    trainable_keras_weights, non_trainable_keras_weights = split_model_weights(
+        keras_model
+    )
+
+    torch_weights_dict: Dict[str, torch.Tensor] = {
+        **trainable_torch_weights,
+        **non_trainable_torch_weights,
+    }
+
+    for keras_weight, keras_weight_name in tqdm(
+        trainable_keras_weights + non_trainable_keras_weights,
+        total=len(trainable_keras_weights + non_trainable_keras_weights),
+        desc="Transferring weights",
+    ):
+        torch_weight_name: str = keras_weight_name
+        for keras_name_part, torch_name_part in weight_name_mapping.items():
+            torch_weight_name = torch_weight_name.replace(
+                keras_name_part, torch_name_part
+            )
+        torch_weight_name = re.sub(
+            r"layerscale\.(\d+)\.variable(?:\.\d+)?", r"gamma_\1", torch_weight_name
+        )
+
+        if "attention" in torch_weight_name:
+            transfer_attention_weights(
+                keras_weight_name,
+                keras_weight,
+                torch_weights_dict,
+                attn_weight_replacement,
+            )
+            continue
+
+        if torch_weight_name not in torch_weights_dict:
+            raise WeightMappingError(keras_weight_name, torch_weight_name)
+
+        torch_weight: torch.Tensor = torch_weights_dict[torch_weight_name]
+
+        if torch_weight_name == "cls_token":
+            keras_weight.assign(torch_weight)
+            continue
+
+        if torch_weight_name == "pos_embed":
+            keras_weight.assign(torch_weight)
+            continue
+
+        if not compare_keras_torch_names(
+            keras_weight_name, keras_weight, torch_weight_name, torch_weight
+        ):
+            raise WeightShapeMismatchError(
+                keras_weight_name,
+                keras_weight.shape,
+                torch_weight_name,
+                torch_weight.shape,
+            )
+
+        transfer_weights(keras_weight_name, keras_weight, torch_weight)
+
+    results = verify_cls_model_equivalence(
+        model_a=torch_model,
+        model_b=keras_model,
+        input_shape=tuple(model_config["input_shape"]),
+        output_specs={"num_classes": model_config["num_classes"]},
+        run_performance=False,
+        atol=1e-4,
+        rtol=1e-4,
+    )
+
+    if not results["standard_input"]:
+        raise ValueError(
+            "Model equivalence test failed - model outputs do not match for standard input"
+        )
+
+    model_filename: str = f"{torch_model_name.replace('.', '_')}.weights.h5"
+    keras_model.save_weights(model_filename)
+    print(f"Model saved successfully as {model_filename}")
+
+    del keras_model, torch_model
+    torch.cuda.empty_cache() if torch.cuda.is_available() else None