keras/hgnetv2_b5_ssld_stage2_ft_in1k

Model Overview

Model Summary

HGNetV2 (Hierarchical Grid Network Version 2) is a modern convolutional neural network (CNN) architecture optimized for GPU efficiency and designed for computer vision tasks. Trained by the model authors on mined ImageNet-22k and ImageNet-1k using SSLD distillation, it excels in edge-focused classification tasks. HGNetV2 delivers strong top-1 accuracy while maintaining a small model footprint, making it ideal for real-time inference on resource-constrained devices.

Key Features:

• GPU-optimized architecture for efficient inference
• Lightweight design with strong performance-to-parameter ratio
• Flexible downsampling strategies (early, late, or no downsampling)
• Modular design with configurable stages and blocks
• Learnable affine transformations for enhanced training stability

Training Strategies:

• SSLD (Semi-Supervised Learning with Distillation): Uses knowledge distillation for improved performance
• Stage 1: Single-stage training approach
• Stage 2: Two-stage training with additional fine-tuning
• ImageNet-22K → ImageNet-1K: Pre-trained on larger dataset, fine-tuned on ImageNet-1K

Weights are released under the Apache 2 License . Keras model code is released under the Apache 2 License.

Links

• HGNetV2 Quickstart Notebook
• HGNetV2 API Documentation
• HGNetV2 Model Card
• KerasHub Beginner Guide
• KerasHub Model Publishing Guide

Installation

Keras and KerasHub can be installed with:

pip install -U -q keras-hub
pip install -U -q keras

Jax, TensorFlow, and Torch come preinstalled in Kaggle Notebooks. For instructions on installing them in another environment see the Keras Getting Started page.

Available HGNetV2 Presets

The following model checkpoints are provided by the Keras team. Full code examples for each are available below.

Preset	Parameters	Description
hgnetv2_b4_ssld_stage2_ft_in1k	13.6M	HGNetV2-B4 model with 2-stage SSLD training, fine-tuned on ImageNet-1K.
hgnetv2_b5_ssld_stage1_in22k_in1k	33.4M	HGNetV2-B5 model with 1-stage SSLD training, pre-trained on ImageNet-22K and fine-tuned on ImageNet-1K.
hgnetv2_b5_ssld_stage2_ft_in1k	33.4M	HGNetV2-B5 model with 2-stage SSLD training, fine-tuned on ImageNet-1K.
hgnetv2_b6_ssld_stage1_in22k_in1k	69.2M	HGNetV2-B6 model with 1-stage SSLD training, pre-trained on ImageNet-22K and fine-tuned on ImageNet-1K.
hgnetv2_b6_ssld_stage2_ft_in1k	69.2M	HGNetV2-B6 model with 2-stage SSLD training, fine-tuned on ImageNet-1K.

Example Usage

# Make necessary imports.
import keras
import keras_hub
import numpy as np

# Load backbone for feature extraction.
backbone = keras_hub.models.HGNetV2Backbone.from_preset(
    "hgnetv2_b5_ssld_stage2_ft_in1k",
)

# Load pre-trained classifier.
classifier = keras_hub.models.HGNetV2ImageClassifier.from_preset(
    "hgnetv2_b5_ssld_stage2_ft_in1k",
)

# Setup parameters.
num_samples = 100
num_classes = 10
image_size = 224

# Use for inference.

# Generate random images - make sure batch size matches.
images = np.random.randint(0, 256, size=(num_samples, image_size, image_size, 3), dtype=np.uint8)

# Generate random labels - same batch size as images.
labels = np.random.randint(0, num_classes, size=(num_samples,), dtype=np.int32)
features = backbone(images)  # Multi-scale features
print(f"Features shape: {features.shape}")
predictions = classifier.predict(images)  # Classification output
print(f"Predictions shape: {predictions.shape}")

# Create custom HGNetV2 backbone.
backbone = keras_hub.models.HGNetV2Backbone(
    depths=[1, 2, 4],
    embedding_size=32,
    hidden_sizes=[64, 128, 256],
    stem_channels=[3, 16, 32],
    hidden_act="relu",
    use_learnable_affine_block=False,
    stackwise_stage_filters=[
    # Stage 0: (in_channels=32, mid_channels=16, out_channels=64, num_blocks=1, num_layers=1, kernel_size=3).
    (32, 16, 64, 1, 1, 3),
    # Stage 1: (in_channels=64, mid_channels=32, out_channels=128, num_blocks=2, num_layers=1, kernel_size=3).
    (64, 32, 128, 2, 1, 3),
    # Stage 2: (in_channels=128, mid_channels=64, out_channels=256, num_blocks=4, num_layers=1, kernel_size=3).
    (128, 64, 256, 4, 1, 3),
    ],
    apply_downsample=[False, True, True],
    use_lightweight_conv_block=[False, False, False],
    image_shape=(224, 224, 3),
)

# Freeze backbone for transfer learning.
classifier.backbone.trainable = False
print("Backbone frozen for transfer learning.")

# Compile for training.
classifier.compile(
    loss=keras.losses.SparseCategoricalCrossentropy(from_logits=True),
    optimizer=keras.optimizers.Adam(5e-5),
    metrics=["accuracy"]
)
print("Classifier compiled.")

# Train on custom dataset.
history = classifier.fit(x=images, y=labels, batch_size=32)
print(f"Training completed. Training history: {history.history}")

Example Usage with Hugging Face URI

# Make necessary imports.
import keras
import keras_hub
import numpy as np

# Load backbone for feature extraction.
backbone = keras_hub.models.HGNetV2Backbone.from_preset(
    "hf://keras/hgnetv2_b5_ssld_stage2_ft_in1k",
)

# Load pre-trained classifier.
classifier = keras_hub.models.HGNetV2ImageClassifier.from_preset(
    "hf://keras/hgnetv2_b5_ssld_stage2_ft_in1k",
)

# Setup parameters.
num_samples = 100
num_classes = 10
image_size = 224

# Use for inference.

# Generate random images - make sure batch size matches.
images = np.random.randint(0, 256, size=(num_samples, image_size, image_size, 3), dtype=np.uint8)

# Generate random labels - same batch size as images.
labels = np.random.randint(0, num_classes, size=(num_samples,), dtype=np.int32)
features = backbone(images)  # Multi-scale features
print(f"Features shape: {features.shape}")
predictions = classifier.predict(images)  # Classification output
print(f"Predictions shape: {predictions.shape}")

# Create custom HGNetV2 backbone.
backbone = keras_hub.models.HGNetV2Backbone(
    depths=[1, 2, 4],
    embedding_size=32,
    hidden_sizes=[64, 128, 256],
    stem_channels=[3, 16, 32],
    hidden_act="relu",
    use_learnable_affine_block=False,
    stackwise_stage_filters=[
    # Stage 0: (in_channels=32, mid_channels=16, out_channels=64, num_blocks=1, num_layers=1, kernel_size=3).
    (32, 16, 64, 1, 1, 3),
    # Stage 1: (in_channels=64, mid_channels=32, out_channels=128, num_blocks=2, num_layers=1, kernel_size=3).
    (64, 32, 128, 2, 1, 3),
    # Stage 2: (in_channels=128, mid_channels=64, out_channels=256, num_blocks=4, num_layers=1, kernel_size=3).
    (128, 64, 256, 4, 1, 3),
    ],
    apply_downsample=[False, True, True],
    use_lightweight_conv_block=[False, False, False],
    image_shape=(224, 224, 3),
)

# Freeze backbone for transfer learning.
classifier.backbone.trainable = False
print("Backbone frozen for transfer learning.")

# Compile for training.
classifier.compile(
    loss=keras.losses.SparseCategoricalCrossentropy(from_logits=True),
    optimizer=keras.optimizers.Adam(5e-5),
    metrics=["accuracy"]
)
print("Classifier compiled.")

# Train on custom dataset.
history = classifier.fit(x=images, y=labels, batch_size=32)
print(f"Training completed. Training history: {history.history}")