Skip to content

Conversation

JyotinderSingh
Copy link
Collaborator

@JyotinderSingh JyotinderSingh commented Aug 20, 2025

Description of the change

This change resolves an issue with loading quantized models from presets. Previously, the model's serialized DTypePolicyMap was not correctly passed to the backbone during loading, which caused failures during initialization of quantized layers.

The fix introduces a new _resolve_dtype utility function that determines the correct dtype for the model based on the following rules:

  1. User-specified dtype: If a user explicitly provides a dtype in the from_preset call (e.g., from_preset("bert_tiny_en_uncased", num_classes=2, dtype="float32")), that value is used.

  2. Float type casting: If no user dtype is provided and the saved dtype is a floating-point type (e.g., "float32"), the model will be loaded using the current Keras default dtype policy. This allows for safe casting between different floating-point precisions.

  3. DTypePolicyMap: If no user dtype is provided and the saved dtype is a complex object (like a DTypePolicyMap for quantization), the saved type is used as is. This ensures that quantization configurations are preserved during loading.

Colab Notebook

  1. Failing Case: This notebook demonstrates the original bug where loading the preset fails.
  2. Successful Case: This notebook shows the corrected behavior where the quantized preset loads successfully with this change.

Checklist

  • I have added all the necessary unit tests for my change.
  • I have verified that my change does not break existing code and works with all backends (TensorFlow, JAX, and PyTorch).
  • My PR is based on the latest changes of the main branch (if unsure, rebase the code).
  • I have followed the Keras Hub Model contribution guidelines in making these changes.
  • I have followed the Keras Hub API design guidelines in making these changes.
  • I have signed the Contributor License Agreement.

@JyotinderSingh JyotinderSingh force-pushed the save_quantized_presets branch from fd28a15 to 1b07517 Compare August 20, 2025 06:32
@JyotinderSingh JyotinderSingh marked this pull request as ready for review August 20, 2025 08:14
Copy link
Member

@mattdangerw mattdangerw left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Thanks!

@JyotinderSingh JyotinderSingh changed the title Fixes issue with loading quantized models from presets Support Loading Quantized Models with from_preset() Aug 21, 2025
@JyotinderSingh JyotinderSingh force-pushed the save_quantized_presets branch from 88e2cec to 430d7b9 Compare August 21, 2025 08:24
Copy link
Collaborator Author

@JyotinderSingh JyotinderSingh left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Resolved comments

@divyashreepathihalli divyashreepathihalli added the kokoro:force-run Runs Tests on GPU label Aug 25, 2025
@kokoro-team kokoro-team removed the kokoro:force-run Runs Tests on GPU label Aug 25, 2025
@divyashreepathihalli
Copy link
Collaborator

/gemini review

Copy link

@gemini-code-assist gemini-code-assist bot left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Code Review

This pull request effectively addresses an issue with loading quantized models from presets by introducing a _resolve_dtype utility function and ensuring dtype policies are correctly serialized. The changes are logical and well-tested. I have a couple of minor suggestions to fix a test assertion message and improve docstring formatting to align with the style guide.

Copy link
Member

@mattdangerw mattdangerw left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

lgtm! just a couple nits

@JyotinderSingh JyotinderSingh force-pushed the save_quantized_presets branch from 0161fb9 to 7eb8f1e Compare August 26, 2025 04:17
Copy link
Member

@mattdangerw mattdangerw left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Thanks!

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Labels
None yet
Projects
None yet
Development

Successfully merging this pull request may close these issues.

5 participants