+1 (415) 366-1364

Transfer Learning for Vision on TensorFlow: KerasHub Backbones, Augmentation, and Export

Almost every computer-vision project a client brings us starts the same way: a few thousand labelled images, a deadline, and a question about whether they need to train a model "from scratch." They do not. In 2026 the right move for image classification, and for most detection and segmentation work, is to take a pretrained backbone, freeze most of it, and fine-tune. This tutorial is the workflow we actually use on client projects with Keras 3 and KerasHub on TensorFlow: build the input pipeline, pick a backbone, run a two-stage fine-tune, evaluate honestly, and export something you can deploy.

Setup

pip install -U keras keras-hub tensorflow

KerasHub absorbed KerasCV, so the vision presets now live alongside the language ones under keras_hub.models. The same code runs on the TensorFlow, JAX, or PyTorch backend - see our Keras 3 multi-backend tutorial - but we keep the TensorFlow backend here because the export path to TF Serving and LiteRT is shortest.

import os
os.environ["KERAS_BACKEND"] = "tensorflow"

import keras
import keras_hub
import tensorflow as tf

IMG_SIZE = 224
BATCH = 32

Step 1: an input pipeline that does not starve the GPU

Vision fine-tuning is the classic input-bound workload: JPEG decode and resize on the CPU, matrix multiplies on the GPU. Build the pipeline properly from the start.

AUTOTUNE = tf.data.AUTOTUNE

train_ds = keras.utils.image_dataset_from_directory(
    "data/train", image_size=(IMG_SIZE, IMG_SIZE), batch_size=BATCH,
    label_mode="categorical", shuffle=True, seed=1337,
)
val_ds = keras.utils.image_dataset_from_directory(
    "data/val", image_size=(IMG_SIZE, IMG_SIZE), batch_size=BATCH,
    label_mode="categorical", shuffle=False,
)

class_names = train_ds.class_names
num_classes = len(class_names)

train_ds = train_ds.cache().prefetch(AUTOTUNE)
val_ds = val_ds.cache().prefetch(AUTOTUNE)

cache() is safe here only if the decoded, resized dataset fits in RAM (roughly n_images * 224 * 224 * 3 bytes). If it does not, drop the cache, or export sharded TFRecords first - the trade-offs are in tf.data Pipelines That Keep Your GPU Busy.

Keep augmentation inside the model, not in tf.data. It then runs on the GPU, it is automatically disabled at inference time, and it travels with the exported artifact:

augment = keras.Sequential([
    keras.layers.RandomFlip("horizontal"),
    keras.layers.RandomRotation(0.05),
    keras.layers.RandomZoom(0.1),
    keras.layers.RandomContrast(0.1),
], name="augment")

Only pick augmentations that preserve the label. Horizontal flips are wrong for text or for anything where handedness matters; heavy colour jitter is wrong for defect detection where colour is the defect. This is the single most common bug we find in client vision code.

Step 2: choose a backbone

backbone = keras_hub.models.Backbone.from_preset("resnet_50_imagenet")
preprocessor = keras_hub.layers.ImageConverter.from_preset(
    "resnet_50_imagenet", image_size=(IMG_SIZE, IMG_SIZE)
)

Rules of thumb we give clients:

  • ResNet-50 / EfficientNet-class CNNs - the default. Fast, small, well understood, fine-tune happily on a few thousand images, and export cleanly to LiteRT for edge deployment.
  • Vision Transformers (ViT) - better accuracy ceiling with more data, and better transfer from large pretraining, but hungrier and slower on CPU inference.
  • SigLIP / CLIP-style image encoders - use when you need embeddings for search or zero-shot labelling rather than a fixed label set. These pair naturally with the retrieval pattern in our two-tower recommender tutorial.
  • Task presets (ImageClassifier, segmentation and detection tasks) wrap a backbone with a head and a preprocessor when you want fewer moving parts.

Check keras_hub.models.Backbone.presets for what is available in your installed version rather than trusting a blog post - including this one.

Step 3: stage one, train the head only

inputs = keras.Input(shape=(IMG_SIZE, IMG_SIZE, 3), dtype="uint8")
x = augment(inputs)
x = preprocessor(x)
backbone.trainable = False              # freeze
x = backbone(x, training=False)         # keep BatchNorm in inference mode
x = keras.layers.GlobalAveragePooling2D()(x)
x = keras.layers.Dropout(0.2)(x)
outputs = keras.layers.Dense(num_classes, activation="softmax")(x)
model = keras.Model(inputs, outputs)

model.compile(
    optimizer=keras.optimizers.Adam(1e-3),
    loss="categorical_crossentropy",
    metrics=["accuracy", keras.metrics.AUC(name="auc")],
)
model.fit(train_ds, validation_data=val_ds, epochs=8)

Two details matter. First, training=False on the frozen backbone keeps its BatchNorm statistics frozen; forgetting this quietly destroys pretrained features and is the reason many "transfer learning doesn't work" tickets land on our desk. Second, the model takes raw uint8 images and does its own preprocessing, so the serving client never has to reproduce a normalization formula.

Step 4: stage two, unfreeze and fine-tune gently

backbone.trainable = True

model.compile(
    optimizer=keras.optimizers.Adam(1e-5),   # ~100x smaller
    loss="categorical_crossentropy",
    metrics=["accuracy", keras.metrics.AUC(name="auc")],
)

callbacks = [
    keras.callbacks.EarlyStopping(monitor="val_auc", mode="max",
                                  patience=4, restore_best_weights=True),
    keras.callbacks.ModelCheckpoint("best.keras", monitor="val_auc",
                                    mode="max", save_best_only=True),
    keras.callbacks.TensorBoard(log_dir="logs"),
]
model.fit(train_ds, validation_data=val_ds, epochs=20, callbacks=callbacks)

If the dataset is small, unfreeze only the last block or two instead of the whole backbone. On a 24 GB GPU, adding keras.mixed_precision.set_global_policy("mixed_float16") before building the model typically buys a 1.5-2x speedup; the caveats are in our mixed precision and XLA post.

Step 5: evaluate like the model will be used

Top-line accuracy hides everything that matters in production. We always report at least:

import numpy as np
from sklearn.metrics import classification_report, confusion_matrix

y_true = np.concatenate([y.numpy() for _, y in val_ds]).argmax(axis=1)
y_prob = model.predict(val_ds)
y_pred = y_prob.argmax(axis=1)

print(classification_report(y_true, y_pred, target_names=class_names))
print(confusion_matrix(y_true, y_pred))

Then look at the confusion matrix by cost, not by count: on an inspection line, one missed defect may be worth a hundred false alarms, which means you should be tuning the decision threshold rather than admiring the argmax. Pull the 50 worst-scoring validation images and look at them - in client work, roughly half of the remaining error is usually label noise, not model capacity.

Step 6: export

model.save("vision_model.keras")                 # training artifact
model.export("export/vision/1")                  # SavedModel for TF Serving

The SavedModel goes straight into the setup from Serving TensorFlow Models in Production. For phones, cameras, and embedded boards, convert the same model to LiteRT and quantize it - see Migrating from TensorFlow Lite to LiteRT and our quantization and pruning guide. Int8 post-training quantization on a fine-tuned ResNet typically costs a fraction of a point of accuracy and cuts the model to a quarter of its size.

A realistic checklist

  • Split by entity, not by file: images of the same product, patient, or shift belong on the same side of the train/val split, or your metrics are fiction.
  • Freeze the split and the seed before you start tuning anything.
  • Log the preset name, image size, and class mapping alongside the artifact; a vision model without its class order is a liability.
  • Watch for drift once deployed - new camera, new lighting, new supplier - using the approach in Catching Data Drift Early with TensorFlow Data Validation.

Most vision projects fail on data hygiene and evaluation, not on architecture. If you have images, labels, and a deadline, our TensorFlow consultants can take this from a notebook to a served endpoint - get in touch and tell us what you are trying to classify.