Most deep learning work in 2026 starts from a checkpoint, not from a blank model. You pick an architecture, download pretrained weights, attach a small head or a prompt, and start training. KerasHub is the keras-team’s answer to that workflow: a library of pretrained vision, text, audio and multimodal models built as honest Keras layers first, so they compose with the rest of the Keras 3 ecosystem rather than living beside it. Where the Hugging Face Transformers stack grew from a modeling library into an ecosystem of its own abstractions, KerasHub took the opposite bet: keep the surface area small, make every pretrained model a keras.Model you can fit(), compile(), save() and inspect with the tools you already know.

The library covers a wide menu — BERT, Gemma, Gemma 3, Llama, Mistral, Qwen, Whisper, SAM, ResNet, RetinaNet, diffusion models — but the genuinely interesting engineering lives underneath the model menu, in roughly a dozen files under keras_hub/src. There you find a three-base class hierarchy (Backbone, Task, CausalLM), a checkpoint registry that resolves four different handle schemes including Kaggle Models and Hugging Face, tokenizer layers that are themselves Keras layers, and a generation system with per-backend compiled fast paths. This post walks that framework layer by layer, reading real code from the current master branch (checked out at commit 3ba52a2, which lands EmbeddingGemma2).

One structural note before we dive in: this is a library of libraries post. KerasHub depends on Keras 3, and Keras 3 dispatches every keras.ops call to TensorFlow, JAX, PyTorch, OpenVINO or NumPy at runtime. That single decision shapes almost every file we will read — most obviously in how generation gets compiled, which we will see in detail below.

The KerasHub skeleton: the public API re-exports two base classes and a preset registry; task models wrap backbones, presets resolve to Kaggle or Hugging Face handles, and tokenizers feed everything.

Reading the overview from left to right:

• Task, L22 is the base class for every end-to-end model: it wraps a Backbone plus a Preprocessor and wires automatic preprocessing into fit(), predict() and evaluate().
• Backbone, L11 is the architecture-only base — hidden states in, hidden states out — with from_preset, LoRA helpers and Hugging Face export.
• CausalLM, L20 extends Task with a sampler argument on compile() and the famous string-in/string-out generate() method.
• get_preset_loader, L553 inspects a preset’s config.json and dispatches to a Keras, Transformers or timm loader — this is how hf:// checkpoints load at all.
• get_file, L133 resolves a preset identifier to a local path across kaggle://, hf://, modelscope://, gs:// and plain directories.
• SentencePieceTokenizer, L32 and BytePairTokenizer, L217 are the two workhorse tokenizer layers.
• GemmaCausalLM, L14 shows what a concrete model contributes on top of the bases: a KV-cache forward pass and a compilable generation step.
• Sampler, L10 is the base class for decoding strategies — greedy, top-k, beam, contrastive, even speculative sampling.

Backbone, Task, CausalLM: The Three-Bases Design

The split between Backbone and Task is the load-bearing idea. A Backbone subclasses keras.Model but is deliberately not tied to any loss: it outputs final hidden states, nothing more. Its constructor builds the layers first — embeddings, transformer stack, final norm — then wires them into a functional keras.Model via keras.Inputs. That functional-model trick matters more than it looks: because all weights are owned by a functional graph, the standard Keras saving, cloning and distribution machinery works without per-model customization. The class even routes around backend quirks in __setattr__ with an explicit torch-backend branch, a small but honest acknowledgment that “write once, run on three backends” has edges.

Task adds the product-facing half. A task owns a backbone and an optional preprocessor, and its preprocess_samples, L64 hook feeds raw strings or images through the preprocessor inside fit() and predict() — set model.preprocessor = None and you drop back to dense tensors. Two other methods deserve attention. First, quantize, L104 overrides Keras quantization to skip preprocessing layers entirely, because tokenizers hold no quantizable weights and are often unbuilt at quantize time, which would otherwise crash. Second, weight bookkeeping is split: load_task_weights, L232 and save_task_weights handle only weights that live outside the backbone (a classification head, for instance), storing them in a separate task.weights.h5 while the backbone’s shard loads independently. summary() is also overridden to print the preprocessor’s tokenizer vocabulary size using a rich-formatted table.

CausalLM then layers generation on top. Its compile, L63 accepts optimizer="auto", loss="auto" defaults — Adam at 2e-5, sparse cross-entropy — plus a sampler argument resolved through a small serialization helper. The heart is generate, L284, which normalizes inputs, preprocesses, distributes tensors, runs a compiled per-batch generation function, and postprocesses back to strings. The compilation itself happens in make_generate_function, L130: on TensorFlow it wraps generate_step in a tf.function with jit_compile; on JAX it goes further, packing sampler variables, trainable and non-trainable variables into an explicit state tuple and running under jax.jit inside a keras.StatelessScope, reassigning only the sampler’s random-seed variables afterwards; on PyTorch it is a simple torch.no_grad() wrapper. Same API, four different execution realities — that is multi-backend in practice, not in marketing.

from_preset: One Registry, Four Handle Schemes

Every base class exposes from_preset(preset, load_weights=True, **kwargs), and the preset machinery in keras_hub/src/utils/preset_utils.py is the most architecturally interesting file in the repo. The registry itself is two module-level dictionaries: BUILTIN_PRESETS, L74 maps short names like "gemma_2b_en" to metadata, and register_presets, L78 fills it from per-model preset tables at import time. Each entry’s crucial field is a kaggle_handle — the short name is just sugar for a Kaggle Models URI.

Resolution happens in get_file, L133. A kaggle://org/model/keras/variant handle goes through kagglehub.model_download, with careful error translation: a 404 becomes a missing-file error while a 403 becomes “if this model is gated, accept its terms on kaggle.com and retry” — the kind of error message you only write after being bitten. An hf://org/model handle calls hf_hub_download; modelscope:// delegates to ModelScope; a bare path is treated as a local directory; and GCS-style URIs ride TensorFlow’s gfile into the local ~/.keras cache.

The loader dispatch is where cross-ecosystem compatibility lives. get_preset_loader reads the preset’s config.json and inspects one key: registered_name means a serialized Keras object (the native format), model_type means a Hugging Face Transformers config, architecture means timm. So TransformersPresetLoader, L50 can construct a KerasHub backbone directly from a HF repo, mapping weights out of the safetensors shards (for background on that format, see the safetensors tour). The Keras loader’s load_task, L705 assembles backbone, preprocessor and task weights, resolving dtype policies so a bfloat16-saved checkpoint casts cleanly to the current default. Saving is symmetric through KerasPresetSaver and a public keras_hub.upload_preset function, which even generates a starter HF model card on first upload. And when you want to leave the ecosystem, export_to_transformers writes HF-compatible output via export_to_safetensors, L402.

Tokenizer Layers That Behave Like Keras Layers

The second quiet idea in KerasHub is that tokenization is a layer, not a function. Tokenizer, L24 subclasses PreprocessingLayer and declares tokenize, detokenize, vocabulary_size and token_to_id — which means a tokenizer can be serialized inside a keras.Model, shipped in a preset, and restored with assets. (If you want a deeper comparison of tokenizer architectures in general, the tokenizers tour covers the fast-Rust design this layer deliberately does not use.)

SentencePieceTokenizer maintains two execution paths behind one API: a TensorFlow Text SentencepieceTokenizer that runs inside the TF graph — essential for saving the tokenizer into a standalone model — and a pure sentencepiece Python path used everywhere else, including JAX and PyTorch. The proto itself is cleverly persisted: set_proto, L167 base64-encodes the model bytes into the config, dodging Keras’s inability to serialize bytestrings, and caches vocabulary metadata eagerly so vocabulary_size() never touches TF at trace time.

BytePairTokenizer reimplements the OpenAI GPT-2 encoder but TF-graph-compatible, including a regex-splitting trick that substitutes a rare character placeholder for unsplittable special tokens — and it ships the Llama3 split pattern as a built-in constant, which is how LlamaTokenizer matches Meta’s reference behavior token-for-token.

From Install to Your First Model

The install is one package plus whichever backend you favor:

pip install keras-hub
# optional but recommended:
pip install kagglehub huggingface_hub sentencepiece

Then a generative model is three lines, exactly as advertised in the class docstrings:

import keras_hub

gemma_lm = keras_hub.models.GemmaCausalLM.from_preset("gemma_2b_en")
gemma_lm.compile(sampler="top_k")
print(gemma_lm.generate("KerasHub is", max_length=64))

from_preset downloads config.json, the tokenizer, and the weights shard from Kaggle, builds a GemmaBackbone, attaches a GemmaCausalLMPreprocessor, and returns a task ready for both generate() and fit(). Fine-tuning on raw strings is a fit() call with a batch size, and parameter-efficient tuning is built in: call enable_lora, L195 with a rank, which freezes the whole backbone and lora-fies only the query/value einsum-dense layers whose names come from default_lora_layer_names, L191. The LoRA factors then save and reload through save_lora_weights/load_lora_weights in a tiny .lora.h5 file, keyed by layer index rather than layer name — names are not unique across a deep stack, and the comment in the code says exactly that.

Gemma and Llama: One Skeleton, Two Personalities

Concrete models then implement the bases. GemmaBackbone builds a tied ReversibleEmbedding, L107 whose output gets multiplied by sqrt(hidden_dim) — the Gemma-specific embedding scaling — then stacks GemmaDecoderBlock, L14 units, each pairing RMSNorm with a CachedGemmaAttention, L17 that supports grouped-query heads, attention logit soft caps, and sliding-window attention on alternating layers. For anyone scaling to multiple accelerators, get_layout_map, L213 returns a keras.distribution.LayoutMap sharding embedding and FFW kernels tensor-parallel while attention outputs go data-parallel — model parallelism as a static layout map, not a rewrite.

The task side is where inference engineering shows. call_with_cache, L174 runs one forward pass against a preallocated key/value cache of shape [batch, layers, 2, max_length, heads, head_dim], updating each layer’s slice at cache_update_index; generate_step, L230 seeds that cache with a full prompt pass, then hands a next() closure to the sampler so decoding strategy and model math stay decoupled. There is even a public score() method for interpretability work, returning per-token logits or losses with an interception hook between layers. Llama follows the same skeleton with its own flavor: LlamaBackbone plus LlamaAttention, L15 grouped-query attention — and if you care about throughput-critical serving rather than in-process generation, dedicated engines like vLLM remain the right tool; KerasHub’s advantage is that the same object you generated with is the one you fine-tuned.

Honest Limits

No source tour is complete without the trade-offs, and KerasHub has real ones.

Preset downloads route through Kaggle. Every built-in short name is a kaggle:// handle underneath, so the first run of from_preset needs network access and, for gated models like Llama, an accepted license on kaggle.com — the 403 path we saw in get_file. Air-gapped environments need a pre-populated cache or a local preset directory, which works but is on you to assemble.

Multi-backend means multi-performance. Four generate paths means four performance profiles. The JAX path re-wraps state into a StatelessScope for jax.jit; the TF path depends on jit_compile surviving Keras version changes; the torch path is the plainest but historically the least optimized. Same code, and benchmarks across backends routinely differ by multiples — pick your backend by measurement, not preference. Server deployments that demand best latency (vLLM, TensorRT-style engines) will beat any of these paths.

The model menu is smaller than Hugging Face’s. The TransformersPresetLoader closes much of the gap — many popular HF checkpoints load directly — but KerasHub’s natively supported architecture list is a few dozen families, not tens of thousands, and weight mapping is one-to-one per architecture: a new HF architecture works only after someone writes its converter.

Keras 3 version coupling. The code checks at runtime whether your Keras supports sharded weight loading, keras.dtype_policies.get existed only after Keras 3.2, and the torch-backend __setattr__ workarounds are pinned to current Keras internals. Upgrading KerasHub usually means upgrading Keras, and occasionally that upgrade is not optional.

None of these are disqualifying — they are the cost of the bet KerasHub makes: one model object that is simultaneously a Keras layer, a Kaggle-downloadable checkpoint, a trainable network on any of three major backends, and an exportable artifact. Read alongside the Keras 3 source and the Transformers library, it is a masterclass in how little framework a model zoo actually needs.

Enjoyed this post? Never miss out on future posts by following us