CIFAR-10 Transfer Learning: Keras 3

Transfer learning reuses a network trained on one large dataset as a feature extractor for a new task. Here the convolutional part of VGG16, trained on 1.3 million ImageNet photographs, is frozen, and only a small classifier on top of it is trained. The benefit should be largest when labelled examples are scarce, which is the usual situation for rare heliophysical events. We therefore train on 1,000, 5,000, and 45,000 CIFAR-10 images and compare with the small CNN of the CIFAR-10 progression trained from scratch on the same images.

The pretrained weights (about 500 MB) are downloaded on first use. Feature extraction takes several minutes on a CPU.

Keras and PyTorch

This notebook repeats the PyTorch version with the Keras 3 API, running on the same PyTorch backend. compile() sets the optimizer and loss, fit() runs the loop over batches and epochs, and callbacks such as early stopping replace hand-written control flow. The data split, model, and evaluation match the PyTorch notebook.

Setup

The libraries used below.

Show code cell source

Hide code cell source

import hashlib
import json
import os

os.environ["KERAS_BACKEND"] = "torch"
os.environ.setdefault("CUBLAS_WORKSPACE_CONFIG", ":4096:8")  # Deterministic cuBLAS on GPUs.

import keras
import matplotlib.pyplot as plt
import numpy as np
import torch
from keras import layers
from sklearn.metrics import (
    ConfusionMatrixDisplay,
    accuracy_score,
    classification_report,
    confusion_matrix,
)
from sklearn.model_selection import train_test_split

assert keras.backend.backend() == "torch"
print(f"Keras {keras.__version__}, backend {keras.backend.backend()}, PyTorch {torch.__version__}")
import time

import pandas as pd

Keras 3.15.1, backend torch, PyTorch 2.13.0

Load CIFAR-10 and split it

The 50,000 official training images are split, with a fixed seed and stratified by class, into 45,000 for training and 5,000 for validation. The 10,000 official test images are used once, at the end.

SEED = 42
EPOCHS = 10  # Reduce to 1 or 2 for a quicker run.
BATCH_SIZE = 256
CLASS_NAMES = ['airplane', 'automobile', 'bird', 'cat', 'deer', 'dog', 'frog', 'horse', 'ship', 'truck']

keras.utils.set_random_seed(SEED)
torch.use_deterministic_algorithms(True, warn_only=True)

(x_development, y_development), (x_test, y_test) = keras.datasets.cifar10.load_data()
y_development = y_development.reshape(-1)
y_test = y_test.reshape(-1)
all_indices = np.arange(len(y_development))
train_indices, validation_indices = train_test_split(
    all_indices,
    test_size=5000,
    random_state=SEED,
    stratify=y_development,
)
split_signature = hashlib.sha256(
    validation_indices.astype("<i8").tobytes()
).hexdigest()[:16]

# Scale pixel values from 0-255 to 0-1.
x_train = x_development[train_indices].astype("float32") / 255.0
y_train = y_development[train_indices].astype(np.int64)
x_validation = x_development[validation_indices].astype("float32") / 255.0
y_validation = y_development[validation_indices].astype(np.int64)
x_test = x_test.astype("float32") / 255.0
y_test = y_test.astype(np.int64)


assert set(train_indices).isdisjoint(validation_indices)
print(
    f"train={len(y_train):,}, validation={len(y_validation):,}, "
    f"test={len(y_test):,}, image shape={x_train.shape[1:]}"
)

train=45,000, validation=5,000, test=10,000, image shape=(32, 32, 3)

Look at the data

One example per class and the number of training images in each. The classes are balanced, so accuracy is a fair summary of performance.

def as_image(sample):
    # Convert one stored sample to (height, width[, channel]) for plotting.
    if sample.ndim == 1:
        sample = sample.reshape(28, 28)
    if sample.ndim == 3 and sample.shape[0] in (1, 3):
        sample = np.transpose(sample, (1, 2, 0))
    return sample.squeeze()


fig = plt.figure(figsize=(12, 3.6))
grid = fig.add_gridspec(2, 10, height_ratios=[1, 1.2])
for label in range(10):
    axis = fig.add_subplot(grid[0, label])
    image = as_image(x_train[np.flatnonzero(y_train == label)[0]])
    axis.imshow(image, cmap="gray" if image.ndim == 2 else None)
    axis.set_title(CLASS_NAMES[label], fontsize=9)
    axis.axis("off")
counts_axis = fig.add_subplot(grid[1, :])
counts_axis.bar(CLASS_NAMES, np.bincount(y_train, minlength=10), color="0.45")
counts_axis.set(ylabel="Training images")
plt.tight_layout()
plt.show()

Training loop

fit() below compiles the model with the Adam optimizer and the cross-entropy loss, trains it, and uses a callback to restore the weights from the epoch with the best validation accuracy.

def fit(model, epochs, learning_rate=1e-3, train=None, validation=None, verbose=True):
    # Train with Adam and restore the epoch with the best validation accuracy.
    model.compile(
        optimizer=keras.optimizers.Adam(learning_rate),
        loss=keras.losses.SparseCategoricalCrossentropy(from_logits=True),
        metrics=["accuracy"],
    )
    keep_best = keras.callbacks.EarlyStopping(
        monitor="val_accuracy", mode="max", patience=epochs, restore_best_weights=True
    )
    history = model.fit(
        *(train or (x_train, y_train)),
        validation_data=validation or (x_validation, y_validation),
        epochs=epochs,
        batch_size=BATCH_SIZE,
        callbacks=[keep_best],
        verbose=2 if verbose else 0,
    ).history
    if verbose:
        best_epoch = int(np.argmax(history["val_accuracy"]))
        print(
            f"restored epoch {best_epoch + 1} "
            f"(validation accuracy {history['val_accuracy'][best_epoch]:.4f})"
        )
    return history


def predict_classes(model, images):
    return model.predict(images, batch_size=BATCH_SIZE, verbose=0).argmax(axis=1)



def plot_history(history, title=None):
    epochs = np.arange(1, len(history["loss"]) + 1)
    best = epochs[int(np.argmax(history["val_accuracy"]))]
    fig, axes = plt.subplots(1, 2, figsize=(10, 3.4))
    for axis, key, label in zip(axes, ("loss", "accuracy"), ("Cross-entropy loss", "Accuracy")):
        axis.plot(epochs, history[key], marker="o", label="training")
        axis.plot(epochs, history[f"val_{key}"], marker="o", label="validation")
        axis.axvline(best, color="0.6", linestyle=":", label="restored epoch")
        axis.set(title=label, xlabel="Epoch")
        axis.grid(alpha=0.25)
    axes[0].legend()
    if title:
        fig.suptitle(title)
    plt.tight_layout()
    plt.show()

Extract frozen VGG16 features

Each image is enlarged to 64×64 pixels, normalized as VGG16 expects, and passed once through the frozen convolutional layers. Averaging the last feature maps gives a 512-number summary per image. Only the small head defined here is trained.

IMAGE_SIZE = 64  # VGG16 was trained on 224-pixel images; 32 pixels leaves a 1x1 map.
HEAD_EPOCHS = 30
SUBSET_SIZES = [1_000, 5_000, 45_000]
backbone = keras.applications.VGG16(
    include_top=False, weights="imagenet", input_shape=(IMAGE_SIZE, IMAGE_SIZE, 3),
    pooling="avg",
)
backbone.trainable = False  # Freeze the pretrained filters.


def extract_features(images):
    # Resize, apply VGG16's own preprocessing, and average-pool its last maps.
    outputs = []
    for start in range(0, len(images), 256):
        batch = keras.ops.image.resize(images[start : start + 256] * 255.0, (IMAGE_SIZE, IMAGE_SIZE))
        batch = keras.applications.vgg16.preprocess_input(batch)
        outputs.append(keras.ops.convert_to_numpy(backbone(batch, training=False)))
    return np.concatenate(outputs).astype(np.float32)


started = time.perf_counter()
features_train = extract_features(x_train)
features_validation = extract_features(x_validation)
features_test = extract_features(x_test)
print(
    f"feature vectors: {features_train.shape[1]} per image; "
    f"extraction took {time.perf_counter() - started:.0f} s"
)


def make_head():
    # The only trainable part of the transfer model.
    return keras.Sequential(
        [
            keras.Input(shape=(512,)),
            layers.Dense(256, activation="relu"),
            layers.Dropout(0.25),
            layers.Dense(10),
        ]
    )

feature vectors: 512 per image; extraction took 50 s

The comparison network

The same small CNN as in the CIFAR-10 progression, trained from raw pixels for 10 epochs.

def make_scratch_network():
    return keras.Sequential(
        [
            keras.Input(shape=(32, 32, 3)),
            layers.Conv2D(16, 3),
            layers.BatchNormalization(),
            layers.LeakyReLU(negative_slope=0.1),
            layers.Conv2D(32, 3, strides=2),
            layers.BatchNormalization(),
            layers.LeakyReLU(negative_slope=0.1),
            layers.Flatten(),
            layers.Dense(100),
            layers.BatchNormalization(),
            layers.LeakyReLU(negative_slope=0.1),
            layers.Dropout(0.5),
            layers.Dense(10),
        ],
        name="cifar10_small",
    )

Train both on growing subsets

For each subset size, both models see the same training images and are scored on the full validation set. The frozen-feature head receives 30 epochs and the scratch CNN 10; this compares two practical training setups, not the effect of pretraining alone. The test set is not used here.

rows, heads = [], {}
for size in sorted({min(size, len(y_train)) for size in SUBSET_SIZES}):
    subset = np.arange(size)  # The training order is already shuffled.
    head = make_head()
    head_history = fit(
        head, HEAD_EPOCHS,
        train=(features_train[subset], y_train[subset]),
        validation=(features_validation, y_validation),
        verbose=False,
    )
    scratch = make_scratch_network()
    scratch_history = fit(
        scratch, EPOCHS, train=(x_train[subset], y_train[subset]), verbose=False
    )
    heads[size] = head
    rows.append(
        {
            "training images": size,
            "frozen VGG16 + head": max(head_history["val_accuracy"]),
            "small CNN from scratch": max(scratch_history["val_accuracy"]),
        }
    )
    print(f"{size:>6,} images: transfer {rows[-1]['frozen VGG16 + head']:.3f}, "
          f"scratch {rows[-1]['small CNN from scratch']:.3f}")

comparison = pd.DataFrame(rows).set_index("training images")
display(comparison.round(3))
fig, ax = plt.subplots(figsize=(6, 3.8))
for column in comparison:
    ax.plot(comparison.index, comparison[column], marker="o", label=column)
ax.set(
    xscale="log", xlabel="Labelled training images", ylabel="Best validation accuracy",
    title="Validation accuracy versus training-set size",
)
ax.legend()
ax.grid(alpha=0.25)
plt.show()

 1,000 images: transfer 0.635, scratch 0.173

 5,000 images: transfer 0.712, scratch 0.183

45,000 images: transfer 0.781, scratch 0.638

	frozen VGG16 + head	small CNN from scratch
training images
1000	0.635	0.173
5000	0.712	0.183
45000	0.781	0.638

Evaluate the transfer model on the test set

The head trained on all 45,000 images is evaluated once.

model = heads[max(heads)]
test_predictions = predict_classes(model, features_test)
test_accuracy = accuracy_score(y_test, test_predictions)
cm = confusion_matrix(y_test, test_predictions, labels=np.arange(10))
print(f"test accuracy: {test_accuracy:.4f}")
print(
    classification_report(
        y_test, test_predictions, labels=np.arange(10),
        target_names=CLASS_NAMES, digits=3, zero_division=0,
    )
)

fig, ax = plt.subplots(figsize=(7, 6))
ConfusionMatrixDisplay(cm, display_labels=CLASS_NAMES).plot(
    ax=ax, colorbar=False, values_format="d", xticks_rotation=45
)
ax.set_title("Test confusion matrix (rows: true class)")
plt.show()

mistakes = np.flatnonzero(test_predictions != y_test)[:12]
fig, axes = plt.subplots(2, 6, figsize=(12, 4.6))
for axis in axes.flat:
    axis.axis("off")
for axis, index in zip(axes.flat, mistakes):
    image = as_image(x_test[index])
    axis.imshow(image, cmap="gray" if image.ndim == 2 else None)
    axis.set_title(f"{CLASS_NAMES[y_test[index]]} as {CLASS_NAMES[test_predictions[index]]}", fontsize=9)
fig.suptitle("First twelve test mistakes (true as predicted)")
plt.tight_layout()
plt.show()

test accuracy: 0.7636
              precision    recall  f1-score   support

    airplane      0.779     0.798     0.788      1000
  automobile      0.884     0.850     0.866      1000
        bird      0.709     0.698     0.704      1000
         cat      0.638     0.587     0.611      1000
        deer      0.676     0.712     0.693      1000
         dog      0.714     0.692     0.703      1000
        frog      0.762     0.811     0.785      1000
       horse      0.761     0.789     0.775      1000
        ship      0.853     0.838     0.846      1000
       truck      0.859     0.861     0.860      1000

    accuracy                          0.764     10000
   macro avg      0.763     0.764     0.763     10000
weighted avg      0.763     0.764     0.763     10000

What the results show

In the saved runs, the frozen VGG16 features give higher validation accuracy at all three training-set sizes. Their advantage is smaller at 45,000 images than at either smaller subset, although the changes are not monotonic in both implementations. The head and scratch network have different training budgets; this is not a controlled measurement of pretraining alone.

Each library uses its own ImageNet VGG16 weights and matching preprocessing, so the PyTorch and Keras scores need not match.

Try it yourself in Keras

Change one choice at a time and keep the data split and evaluation unchanged:

• add a 250-image subset to SUBSET_SIZES to see how far the frozen features carry.

• replace the head with a single linear layer, a logistic regression on the VGG16 features.

• for a larger extension, replace cached-feature training with an end-to-end image model, unfreeze the last VGG16 block, and try a learning rate of 1e-5.