Downloaded model deserialization failure

[r-zip]

Describe the bug

Models downloaded through KerasHub fail to deserialize.

To Reproduce

It is not possible to reproduce this bug in Colab. I suspect there is some library version incompatibility; however, the Colab environment is so complicated that I was unable to replicate it locally.

To reproduce the bug, create a new directory and put the following pyproject.toml into it:

[project]
name = "foo"
version = "0.0.1"
description = "Add your description here"
readme = "README.md"
requires-python = "==3.11.*"
dependencies = [
    "jax==0.4.33",
    "jax-cuda12-pjrt==0.4.33",
    "jax-cuda12-plugin==0.4.33",
    "jaxlib==0.4.33",
    "jupyter>=1.1.1",
    "kaggle==1.6.17",
    "kagglehub==0.3.6",
    "keras==3.8.0",
    "keras-hub==0.18.1",
    "matplotlib>=3.10.0",
    "tensorflow==2.18.0",
    "tensorflow-datasets==4.9.7",
    "tensorflow-hub==0.16.1",
    "tensorflow-io-gcs-filesystem==0.37.1",
    "tensorflow-metadata==1.16.1",
    "tensorflow-probability==0.24.0",
    "tensorflow-text==2.18.1",
    "tf-keras==2.18.0",
    "tf-slim==1.1.0",
    "torch",
    "torchaudio",
    "torchsummary==1.5.1",
    "torchvision",
]

[tool.uv.sources]
torchaudio = { url = "https://download.pytorch.org/whl/cu124/torchaudio-2.5.1%2Bcu124-cp311-cp311-linux_x86_64.whl" }
torch = { url = "https://download.pytorch.org/whl/cu124/torch-2.5.1%2Bcu124-cp311-cp311-linux_x86_64.whl" }
torchvision = { url = "https://download.pytorch.org/whl/cu124/torchvision-0.20.1%2Bcu124-cp311-cp311-linux_x86_64.whl" }

The versions of these packages (besides matplotlib and Jupyter) are pinned to their values in Colab. Next, run uv sync, then source .venv/bin/activate. Finally, copy and paste the following code blocks into a Jupyter notebook (taken from the tutorial here):

import os

os.environ["KERAS_BACKEND"] = "jax"  # or "tensorflow" or "torch"
os.environ["XLA_PYTHON_CLIENT_MEM_FRACTION"] = "1.0"

import keras
import numpy as np
import matplotlib.pyplot as plt

image_url = "https://upload.wikimedia.org/wikipedia/commons/a/aa/California_quail.jpg"
image_path = keras.utils.get_file(origin=image_url)
image = keras.utils.load_img(image_path)
plt.imshow(image)

import keras_hub

image_classifier = keras_hub.models.ImageClassifier.from_preset(
    "resnet_50_imagenet",
    activation="softmax",
)
batch = np.array([image])
image_classifier.preprocessor.image_size = (224, 224)
preds = image_classifier.predict(batch)
preds.shape

The final cell results in the following error:

TypeError: <class 'keras_hub.src.models.resnet.resnet_image_classifier.ResNetImageClassifier'> could not be deserialized properly. Please ensure that components that are Python object instances (layers, models, etc.) returned by `get_config()` are explicitly deserialized in the model's `from_config()` method.

config={'module': 'keras_hub.src.models.resnet.resnet_image_classifier', 'class_name': 'ResNetImageClassifier', 'config': {'backbone': {'module': 'keras_hub.src.models.resnet.resnet_backbone', 'class_name': 'ResNetBackbone', 'config': {'name': 'res_net_backbone', 'trainable': True, 'input_conv_filters': [64], 'input_conv_kernel_sizes': [7], 'stackwise_num_filters': [64, 128, 256, 512], 'stackwise_num_blocks': [3, 4, 6, 3], 'stackwise_num_strides': [1, 2, 2, 2], 'block_type': 'bottleneck_block', 'use_pre_activation': False, 'image_shape': [None, None, 3], 'dtype': None}, 'registered_name': 'keras_hub>ResNetBackbone'}, 'preprocessor': {'module': 'keras_hub.src.models.resnet.resnet_image_classifier_preprocessor', 'class_name': 'ResNetImageClassifierPreprocessor', 'config': {'name': 'res_net_image_classifier_preprocessor', 'trainable': True, 'dtype': {'module': 'keras', 'class_name': 'DTypePolicy', 'config': {'name': 'float32'}, 'registered_name': None}, 'image_converter': {'module': 'keras_hub.src.models.resnet.resnet_image_converter', 'class_name': 'ResNetImageConverter', 'config': {'name': 'res_net_image_converter', 'trainable': True, 'dtype': {'module': 'keras', 'class_name': 'DTypePolicy', 'config': {'name': 'float32'}, 'registered_name': None}, 'image_size': [224, 224], 'scale': [0.017124753831663668, 0.01750700280112045, 0.017429193899782133], 'offset': [-2.1179039301310043, -2.0357142857142856, -1.8044444444444445], 'interpolation': 'bicubic', 'crop_to_aspect_ratio': True}, 'registered_name': 'keras_hub>ResNetImageConverter'}, 'config_file': 'preprocessor.json'}, 'registered_name': 'keras_hub>ResNetImageClassifierPreprocessor'}, 'name': 'res_net_image_classifier', 'num_classes': 1000, 'pooling': 'avg', 'activation': 'softmax', 'dropout': 0.0}, 'registered_name': 'keras_hub>ResNetImageClassifier'}.

Exception encountered: <class 'keras_hub.src.models.resnet.resnet_backbone.ResNetBackbone'> could not be deserialized properly. Please ensure that components that are Python object instances (layers, models, etc.) returned by `get_config()` are explicitly deserialized in the model's `from_config()` method.

config={'module': 'keras_hub.src.models.resnet.resnet_backbone', 'class_name': 'ResNetBackbone', 'config': {'name': 'res_net_backbone', 'trainable': True, 'input_conv_filters': [64], 'input_conv_kernel_sizes': [7], 'stackwise_num_filters': [64, 128, 256, 512], 'stackwise_num_blocks': [3, 4, 6, 3], 'stackwise_num_strides': [1, 2, 2, 2], 'block_type': 'bottleneck_block', 'use_pre_activation': False, 'image_shape': [None, None, 3], 'dtype': None}, 'registered_name': 'keras_hub>ResNetBackbone'}.

Exception encountered: unsupported operand type(s) for %: 'NoneType' and 'int'

Expected behavior

I expect the last cell to produce the output

 1/1 ━━━━━━━━━━━━━━━━━━━━ 2s 2s/step
(1, 1000)

from the tutorial page.

Additional context

The tutorial works on Colab, but not locally; I tried this both on Linux with CUDA GPUs and on my M1 MacBook Pro. Neither worked.

Would you like to help us fix it?
I don't believe I have the expertise or time to fix it, unfortunately.

[mehtamansi29]

Hi @r-zip -

Thanks for reporting the issue. I reproduce the issue in keras_hub(0.18.1) and keras(3.8.0) version and issue doesn't reproduce anymore. Attached gist here for the reference.

[r-zip]

Hi, yes, I mentioned this when I wrote the issue. However, when installing these packages locally (both on MacOS and on Linux), the example breaks for me. Could you please try reproducing the bug outside of the Colab environment?

[mehtamansi29]

Hi @r-zip -

I tried to reproduce the same issue outside colab environment(vscode in MAC os) and it is working fine for me.

[yashwanth510]

Hey team! 👋

I've made a PR to fix this.

Cause: head_dtype was passed to __init__() and used to set the dtype policy on pooler, output_dropout, and output_dense, but it was never saved on self or included in get_config(), so it was lost every time the program was saved and loaded.
Fix: In image_classifier.py, change 2 lines: 1. self.head_dtype = head_dtype in __init__(), and 2. "head_dtype": self.head_dtype in get_config().

All local verification checks pass ✅

PR: #2614