feat: add Keras 3 support for multi-backend (pytorch, jax)#5650
Open
rmarquis wants to merge 8 commits into
Open
feat: add Keras 3 support for multi-backend (pytorch, jax)#5650rmarquis wants to merge 8 commits into
rmarquis wants to merge 8 commits into
Conversation
…ack) Make TensorFlow a lazy dependency and branch save/load/runnable logic based on keras.config.backend(). This allows bentoml.keras to work with Keras models using the PyTorch backend without requiring TensorFlow. - Detect active backend via keras.config.backend() (Keras 3) or default to tensorflow (Keras 2). - Record backend-relevant versions in ModelContext (keras, torch, jax). - Reject tf_signatures/tf_save_options when backend is not tensorflow. - Add backend-aware device handling and tensor conversion in runnable. - Preserve existing TensorFlow backend behavior.
- Add tests/integration/frameworks/test_keras_torch_backend.py covering save/load/runner inference for a Keras 3 model using the torch backend. - Add a note to docs/source/reference/bentoml/frameworks/keras.rst about Keras 2/3 multi-backend support and TensorFlow-only options.
For more information, see https://pre-commit.ci
There was a problem hiding this comment.
Pull request overview
This PR extends bentoml.keras to support Keras 3 multi-backend setups (TensorFlow / PyTorch / JAX) by making TensorFlow optional where possible and adding backend-specific runner behavior.
Changes:
- Add backend detection and backend-specific runnable execution paths (TensorFlow vs. torch vs. other backends).
- Adjust save/load behavior and recorded framework context for Keras 3 usage.
- Add an integration test for saving/loading with the Keras torch backend and update reference docs.
Reviewed changes
Copilot reviewed 3 out of 3 changed files in this pull request and generated 5 comments.
| File | Description |
|---|---|
| tests/integration/frameworks/test_keras_torch_backend.py | Adds an integration test covering save/load + runner execution under Keras torch backend. |
| src/bentoml/_internal/frameworks/keras.py | Implements backend-aware loading and runnable execution, including torch/jax tensor conversion helpers. |
| docs/source/reference/bentoml/frameworks/keras.rst | Documents multi-backend support and TensorFlow-only options. |
💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.
Wire _normalize_device_name() into load_model() and KerasRunnable so TensorFlow-style device strings (e.g. '/device:CPU:0') are accepted when the active Keras backend is PyTorch. Rename the previously-unused _normalize_torch_device helper to _normalize_device_name and extend its docstring. Add test_keras_torch_load_with_tf_style_device to verify save/load and runner inference with a TF-style device string under the torch backend.
Store the active Keras backend name in ModelContext.framework_versions under the 'backend' key. In load_model(), compare it against the current backend from keras.config.backend() and raise BentoMLException with a helpful message when they differ. This prevents confusing failures when a model saved with one backend is loaded under another. Update the load_model docstring to accurately describe this behavior. Add test_keras_torch_load_with_wrong_backend_raises to verify the mismatch error.
Avoid IndexError when a TensorFlow-style GPU device name is requested
but tf.config.list_physical_devices('GPU') returns an empty list. Only
call set_memory_growth when at least one physical GPU is present.
The torch backend runnable only converted numpy arrays and torch.Tensor inputs, leaving plain Python lists/tuples to pass through unchanged. Keras with the torch backend can fail on raw sequences, so convert them with torch.as_tensor() and move to the target device. Add test_keras_torch_runner_accepts_list_input covering list and tuple inputs through the runner.
rmarquis
changed the title
The Keras 3 integration only handles tensorflow, torch, and jax in _get_context. The save_model docstring and framework reference docs still listed OpenVINO as a supported backend, so drop it to match the actual implementation.
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
What does this PR address?
Fixes #5649
Before submitting:
guide on how to create a pull request.
pre-commit run -ascript has passed (instructions)?those accordingly? Here are documentation guidelines and tips on writting docs.