🗺️ Keras Development Roadmap

[fchollet]

Here's an overview of the features we intend to work on in the near future, across Core Keras, KerasNLP, and KerasCV.

Core Keras

Saving & export

• [Open for Contributions] Add utility for saved Keras model file inspection, diff, and patching.
• Implement saving support for sharded models (sharded weights files).
• Improve model export story:
  • [Open for Contributions] Implement export to ONNX (from any backend)
  • [Open for Contributions] Implement export to TensorRT (from any backend)
  • [Open for Contributions] Write serving/export guide explaining the different servings options (SavedModel, ONNX, TensorRT, TorchServe, etc.)

Distribution

• Add PyTorch SPMD support in keras.distribution.
• Add TensorFlow SPMD (DTensor) support in keras.distribution.

Performance

• Enable PyTorch compilation on any Keras model.
• Also make it possible to provide torch compilation arguments to Keras.
• [Open for Contributions] Add a profiler Callback. With any backend, this callback should be able to give you a report on the computation/time breakdown of your layers and your train step.
• [Open for Contributions] Add ops.flash_attention ops for Pytorch and JAX (implemented via Pallas).
• ✅ Add Float8 training.
• ✅ Add LoRA support to key layer types (Dense, EinsumDense, Conv layers)
• ✅ Enable QLoRA
• ✅ Enable int8 inference for key layer types (Dense, EinsumDense)
• ✅ Enable float8 training for key layer types (Dense, EinsumDense)

Modeling

• ✅ Bring back deeply nested Functional inputs/outputs. Right now it is enforced that the inputs and outputs argument in Model should be flat structured (e.g. flat dicts or lists). However it may be useful to support e.g. dicts of lists, or dicts of dicts.
• ✅ Add support for optional inputs in Functional models.
• ✅ Add multi-backend scan operation in the style of JAX.
• [Open for Contributions] Increase coverage of NumPy API in keras.ops.
  • Any function from the NumPy API is in scope for inclusion. If you notice a missing op that you need, please open a PR!
  • To prioritize:
    • ✅ ops.argpartition
    • ✅ ops.vectorize
    • ✅ ops.select
    • ✅ ops.eigh
• Add tensor.at operator.
• Add sparse support to PyTorch backend.
• Increase coverage of sparse support in keras.ops:
  • Reduction ops
  • Boolean ops

Ecosystem

• Add integration with Hugging Face to enable one-line upload/download of .keras models.
• [Open for Contributions] Finish implementing the MLX backend, in collaboration with Apple (see mlx branch).
• Add OpenVINO inference-only backend, in collaboration with Intel.

Code health

• ✅ Check in the output of pip_build.py to make the Keras package runnable from its sources without any postprocessing.

Keras "starter packs"

• [Open for Contributions] Create a Recommender System "starter packer": a standalone repository containing concise, readable, customizable code to implement a recommender system with Keras 3 (supporting all backends), including model code, data loading code, training code, and model export for serving.
• [Open for Contributions] Large-scale LLM training starter pack using keras.distribution API and JAX backend.

Keras Models

• Unify KerasCV and KerasNLP into a single Keras Models package.

KerasCV

• ✅ Add an integration with Hugging Face to enable one-line upload of fine-tuned models to Hugging Face Hub.
• Make StableDiffusion model conform to from_preset API.
• Add support for channels_first format in all models. Necessary for PyTorch performance.
• Also add channels_first testing on CI.
• Add standard image classifiers to KerasCV
  • All Keras Applications models should be brought to KerasCV.
  • Add more pretrained weights to standard image classifiers, in particular OpenImages.
    • Right now we only have ImageNet
• Drop Keras 2 support.
• [Open for Contributions] Scrub TF references from docstrings.
• Scrub Waymo-specific C++ code.
• [Open for Contributions] Add new models:
  • OCR
  • CoCa
  • Faster RCNN
  • YOLOv8 Instance segmentation
  • VVIT
  • Grounding Dino
  • Visual Q&A task model with CLIP
  • StableDiffusion v3
  • Mask RCNN
  • YOLOv9
  • Depth anything
  • layoutlmv3

KerasNLP

• ✅ Add an integration with Hugging Face to enable one-line upload of fine-tuned models to Hugging Face Hub.
• Remove dependency on tf-text for KerasNLP. Reimplement in Python.
• Add support for dynamic sequence length inference for PyTorch LLMs in KerasNLP, to bridge the performance gap with Hugging Face.
• Add support for QLoRA to select KerasNLP models.
• Drop Keras 2 support.
• Add FlashAttention support for select models.
• [Open for Contributions] Add new models:
  • Llama2, 3
  • Falcon
  • Mixtral
  • Bloom
  • Flan
  • MPT
  • Vicuna

[pure-rgb]

Unify KerasCV and KerasNLP into a single Keras Models package.

1. does it mean that kerascv and kerasnlp will be merged together into one single package, like tensorflow model garden?

2. kerascv development is terribly slow, and lacks of resources. no release since february (0.8.2), may models doesn't have weight, i.e. vit, mit. Object detection models that is added (yolo-v8, retina-net, faster-rcnn) all are broken and doesn't have evaluation scripts to holds their validity. Generative model only stable diffusion, where many (too many impactfull models are available now).

3. like huggigface, why enigneer or researcher are not hired and reinformce into this package. Check this repo, the ownver can add any new models written in keras. so why not official keras team? Is it because the api desgin is too rich and complicated for contributor? Not to mention, many reviews are pending, contributor left at th end.

[IMvision12]

I concur with @pure-rgb's observation that the API design has become overly elaborate and intricate for contributors. Previously, it was straightforward and comprehensible, but with the introduction of the Backbone API, it has become more convoluted.

[pure-rgb]

 with the introduction of the Backbone API, it has become more convoluted.

The person who introduced such API design already left google. :p

[kernel-loophole]

@fchollet I've been going through the features you're planning to implement, and I'm particularly interested in contributing to KerasNLP. Specifically, I'm eager to get involved in the development of dynamic sequence length inference for PyTorch LLMs.

[fchollet]

@kernel-loophole please open an issue on the KerasNLP repo if you'd like to contribute this feature!

[kernel-loophole]

okay thanks

[MoFHeka]

I am a developer of tensorflow recommenders-addons and I now need to develop an all-to-all embedding layer for multi-GPU distributed training of recommendation models. The old tensorflow distributed strategy clearly did not meet this need. So the question is, should I develop on TF DTensor or Jax? Because it seems that Keras support for TF DTensor is not friendly. But Jax lacks the ability to online inference services and the functional components used by various recommendation algorithms. Also recommenders-addons has a lot of custom operators.

[jeffcarp]

online inference services and the functional components used by various recommendation algorithms

@MoFHeka Can you elaborate on what you need here?

[MoFHeka]

online inference services and the functional components used by various recommendation algorithms

@MoFHeka Can you elaborate on what you need here?

@jeffcarp If a third-party custom op (primitives) is used in Jax training, it will be difficult to convert it to saved_model for online inference. Jax2tf is not easy to use. Unlike TF custom op, it only needs to compile with TFServing or preload the dynamic link library.

You may not be very concerned about how DTensor or Jax will evolve in the future. But for now the large number of recommendation models are trained by Keras, I'm interested to hear what I think of these two frameworks as a Keras developer. After all, both frameworks have their own problems for us. Since one of our priorities is to be compatible with Keras APIs, I would like to know whether Keras currently prefers to do more development for DTensor or Jax and integrate their which features? Or which of the two is more worthy of support?