fix: serialize head_dtype in all ImageClassifier subclasses

[yashwanth510]

Problem

Turns out, the head_dtype serialization bug from #2614 pops up in six other ImageClassifier subclasses too.

Here’s what happened: head_dtype gets passed into __init__() and helps set the dtype policy for classifier head layers. But nobody actually saved it on self or included it in get_config(). So, every time you saved or loaded the model, that value just disappeared without warning.

Affected Models

• VitImageClassifier → vit/vit_image_classifier.py
• DeiTImageClassifier → deit/deit_image_classifier.py
• VGGImageClassifier → vgg/vgg_image_classifier.py
• MobileNetImageClassifier → mobilenet/mobilenet_image_classifier.py
• MobileNetV5ImageClassifier → mobilenetv5/mobilenetv5_image_classifier.py
• HGNetV2ImageClassifier → hgnetv2/hgnetv2_image_classifier.py

Fix

Just like in #2614, I fixed all six subclasses with two quick changes:

1. Saved head_dtype in __init__():
   self.head_dtype = head_dtype

2. Added head_dtype to get_config():
   "head_dtype": self.head_dtype,

Result

• Now head_dtype sticks around when saving and loading models.
• Serialization matches the base ImageClassifier.
• No breaking changes.

Related

This builds on #2614, which patched the same problem in the base ImageClassifier.

cc: @mattdangerw @sachinprasadhs @divyashreepathihalli

[gemini-code-assist]

Summary of Changes

Hello, I'm Gemini Code Assist¹! I'm currently reviewing this pull request and will post my feedback shortly. In the meantime, here's a summary to help you and other reviewers quickly get up to speed!

This pull request resolves a widespread serialization issue across multiple ImageClassifier subclasses where the head_dtype configuration was not being preserved when models were saved and reloaded. By correctly storing and exposing this parameter for serialization, the change guarantees the integrity of model configurations and ensures consistent behavior after persistence, preventing silent data type policy changes in classifier heads.

Highlights

• Serialization Bug Fix: Addressed a critical serialization bug where the head_dtype parameter was not being saved or loaded correctly in several ImageClassifier subclasses, leading to data loss upon model persistence.
• Affected Models: The fix has been applied to VitImageClassifier, DeiTImageClassifier, VGGImageClassifier, MobileNetImageClassifier, MobileNetV5ImageClassifier, and HGNetV2ImageClassifier.
• Implementation Details: The head_dtype is now explicitly saved as an instance variable (self.head_dtype) within the __init__ method and included in the get_config() method for proper serialization.
• Consistency and Compatibility: Ensures serialization consistency with the base ImageClassifier class and introduces no breaking changes, building upon a previous fix in The solution is to serialize head_dtype in `ImageClassifier.get_config() #2614.

🧠 New Feature in Public Preview: You can now enable Memory to help Gemini Code Assist learn from your team's feedback. This makes future code reviews more consistent and personalized to your project's style. Click here to enable Memory in your admin console.

Changelog

• keras_hub/src/models/deit/deit_image_classifier.py
  • Stored head_dtype as an instance variable.
  • Included head_dtype in the get_config method for serialization.
• keras_hub/src/models/hgnetv2/hgnetv2_image_classifier.py
  • Stored head_dtype as an instance variable.
  • Included head_dtype in the get_config method for serialization.
• keras_hub/src/models/mobilenet/mobilenet_image_classifier.py
  • Stored head_dtype as an instance variable.
  • Included head_dtype in the get_config method for serialization.
• keras_hub/src/models/mobilenetv5/mobilenetv5_image_classifier.py
  • Stored head_dtype as an instance variable.
  • Included head_dtype in the get_config method for serialization.
• keras_hub/src/models/vgg/vgg_image_classifier.py
  • Stored head_dtype as an instance variable.
  • Included head_dtype in the get_config method for serialization.
• keras_hub/src/models/vit/vit_image_classifier.py
  • Stored head_dtype as an instance variable.
  • Included head_dtype in the get_config method for serialization.

Activity

• This pull request builds on the work done in The solution is to serialize head_dtype in `ImageClassifier.get_config() #2614, which addressed a similar head_dtype serialization issue in the base ImageClassifier.
• The author has proactively identified and addressed the recurrence of this bug across six additional ImageClassifier subclasses.
• Potential reviewers @mattdangerw, @sachinprasadhs, and @divyashreepathihalli have been CC'd for their attention.

Using Gemini Code Assist

The full guide for Gemini Code Assist can be found on our documentation page, here are some quick tips.

Invoking Gemini

You can request assistance from Gemini at any point by creating a comment using either /gemini <command> or @gemini-code-assist <command>. Below is a summary of the supported commands on the current page.

Feature	Command	Description
Code Review	/gemini review	Performs a code review for the current pull request in its current state.
Pull Request Summary	/gemini summary	Provides a summary of the current pull request in its current state.
Comment	@gemini-code-assist	Responds in comments when explicitly tagged, both in pull request comments and review comments.
Help	/gemini help	Displays a list of available commands.

Customization

To customize Gemini Code Assist for GitHub experience, repository maintainers can create a configuration file and/or provide a custom code review style guide (such as PEP-8 for Python) by creating and adding files to a .gemini/ folder in the base of the repository. Detailed instructions can be found here.

Limitations & Feedback

Gemini Code Assist may make mistakes. Please leave feedback on any instances where its feedback is incorrect or counter productive. You can react with 👍 and 👎 on @gemini-code-assist comments. If you're interested in giving your feedback about your experience with Gemini Code Assist for Github and other Google products, sign up here.

You can also get AI-powered code generation, chat, as well as code reviews directly in the IDE at no cost with the Gemini Code Assist IDE Extension.

Footnotes

1. Review the Privacy Notices, Generative AI Prohibited Use Policy, Terms of Service, and learn how to configure Gemini Code Assist in GitHub here. Gemini can make mistakes, so double check it and use code with caution. ↩

[gemini-code-assist]

Code Review

This pull request effectively addresses a critical serialization bug across multiple ImageClassifier subclasses. By correctly saving the head_dtype attribute in the __init__ method and including it in the get_config() method, the changes ensure that the model's dtype policy for the classification head is preserved during serialization and deserialization. This is crucial for maintaining model correctness and consistency, aligning with the repository's layer implementation guidelines for proper attribute handling and configuration serialization.

[sachinprasadhs]

Thanks!

[sachinprasadhs]

Fix the failing error and you can include the changes in #2614 in this PR.

[yashwanth510]

Thanks for the review, @sachinprasadhs!

I'll put self.head_dtype under # === Config === in all the files. I'll also add the fix for image_classifier.py from #2614 to this PR.

Should I close #2614 after this is fixed?

[yashwanth510]

I made the requested changes, @sachinprasadhs, by

• moving self.head_dtype under # === Config ===
  (the config section) in all 7 files

• and including the base ImageClassifier fix (from The solution is to serialize head_dtype in `ImageClassifier.get_config() #2614) in this pull request.

Please let me know if you need any more changes!

[sachinprasadhs]

You need to fix the failing test cases

[yashwanth510]

Thank you for the review.

I have fixed the issues. Now all the test cases are passing.

[yashwanth510]

Hi @sachinprasadhs I fixed the tests that were failing. Now all checks are passing. Can you please take a look when you have time? Thanks!

[sachinprasadhs]

Thanks, but you still need to handle deserializing when it is loaded from config.
It needs to be done at the root level.
Also, add a test case to serialize and save model config and load the saved config.

[yashwanth510]

I appreciate your input, @sachinprasadhs.

In order to ensure that head_dtype is preserved, I'll add deserialization handling for it at the root level in ImageClassifier along with a test case that saves and reloads the model configuration.

I'll release the updates soon. If there is anything else you would like me to change, do let me know.

[sachinprasadhs]

root level in the sense task.py, not the ImageClassifier, test case also goes there.

[divyashreepathihalli]

/gemini review

[gemini-code-assist]

Code Review

This pull request correctly addresses a serialization bug by ensuring head_dtype is properly saved and retrieved in get_config across several ImageClassifier subclasses. The changes are consistent and effectively fix the issue. I've added a few suggestions to refactor some of the get_config methods to reduce code redundancy and improve maintainability by better leveraging inheritance.

[yashwanth510]

Thanks for clarifying @sachinprasadhs!

I'll add the head_dtype deserialization handling in task.py
(not ImageClassifier) and add the test case there as well.

Will push the update shortly!

[yashwanth510]

Hello @sachinprasadhs, I've made the modifications you asked for:

1. Added the root-level head_dtype deserialization in task.py and from_config().
2. The test case "test_image_classifier_head_dtype_serialization" was added to "task_test.py."

Every test is successful locally. Tell me if anything else needs to be changed, please!