TensorFlow CNNs and Transfer Learning

Real-World Scenario

A plant nursery wants an app that tells customers which flower they photographed. There are only a few thousand labelled photos — too few to train a large image model from scratch, but plenty to adapt one that has already learned to see from ImageNet. This lesson builds up convolutions from their shapes, trains a small CNN from scratch, then does exactly that adaptation with a pretrained MobileNetV2.

All examples ran with TensorFlow 2.21 and Keras 3.15 on a CPU-only machine. TensorFlow’s startup log lines on stderr (no CUDA driver found, oneDNN enabled) are left out of the outputs below.

Convolution Output Shapes

Keras images are channels-last: (batch, height, width, channels). Padding and strides decide the output’s height and width; the number of filters decides its channels.

import numpy as np
import keras

x = np.random.rand(8, 28, 28, 1).astype("float32")   # (batch, height, width, channels)

same = keras.layers.Conv2D(32, 3, padding="same")
valid = keras.layers.Conv2D(32, 3, padding="valid")
strided = keras.layers.Conv2D(32, 3, padding="same", strides=2)
pool = keras.layers.MaxPooling2D(2)

print("same:   ", same(x).shape)
print("valid:  ", valid(x).shape)
print("strides:", strided(x).shape)
print("pooled: ", pool(same(x)).shape)
print("kernel: ", same.kernel.shape)

$ python conv_shapes.py
same:    (8, 28, 28, 32)
valid:   (8, 26, 26, 32)
strides: (8, 14, 14, 32)
pooled:  (8, 14, 14, 32)
kernel:  (3, 3, 1, 32)

"same" pads the border so the output stays 28×28; "valid" — the Keras default — pads nothing, so a 3×3 window fits only 26 positions across. A stride of 2 and a 2×2 max-pool both halve the resolution to 14×14. The kernel is (height, width, in_channels, filters): 32 filters, each 3×3 over one input channel. PyTorch orders both the images and the kernel differently — channels first — which matters when porting weights.

Weight Sharing

Both layers below map a 28×28 image to the same 26×26×32 output. Predict which needs more parameters, and by how much.

import keras

conv = keras.layers.Conv2D(32, 3)
dense = keras.layers.Dense(26 * 26 * 32)       # same input/output sizes, no sharing
conv.build((None, 28, 28, 1))
dense.build((None, 28 * 28))

print(f"conv:  {conv.count_params():>12,}")
print(f"dense: {dense.count_params():>12,}")

$ python conv_params.py
conv:           320
dense:   16,981,120

The convolution needs 3 × 3 × 1 × 32 weights plus 32 biases — 320 numbers — because each small filter is reused at every position. The dense layer learns a separate weight for every input-output pair: 784 × 21,632 + 21,632, almost 17 million. That reuse is why CNNs are cheap, and why a pattern learned in one corner is detected everywhere.

What a Filter Computes

A filter is a small grid of weights; the output at each position is the sum of the weights times the pixels underneath. Set one by hand to see it.

import numpy as np
import keras

image = np.zeros((1, 5, 6, 1), dtype="float32")
image[:, :, 3:, :] = 1.0                 # dark left half, bright right half
print(image[0, :, :, 0])

conv = keras.layers.Conv2D(1, 3, use_bias=False)
conv.build(image.shape)
kernel = np.array([[-1, 0, 1],
                   [-1, 0, 1],
                   [-1, 0, 1]], dtype="float32")   # right minus left
conv.set_weights([kernel.reshape(3, 3, 1, 1)])
print(conv(image).numpy()[0, :, :, 0])

$ python edge_filter.py
[[0. 0. 0. 1. 1. 1.]
 [0. 0. 0. 1. 1. 1.]
 [0. 0. 0. 1. 1. 1.]
 [0. 0. 0. 1. 1. 1.]
 [0. 0. 0. 1. 1. 1.]]
[[0. 3. 3. 0.]
 [0. 3. 3. 0.]
 [0. 3. 3. 0.]]

The kernel subtracts each window’s left column from its right column. Over a flat region — all dark or all bright — that is zero. Where the window straddles the dark-to-bright boundary, three bright pixels on the right minus three dark on the left gives 3. The output is 3×4 because Conv2D defaults to "valid" padding. Training learns kernels like this from data instead of having them written by hand.

Training a CNN on Fashion-MNIST

Fashion-MNIST has the same shape as MNIST — 28×28 grayscale, 10 classes — but shows clothing, which is harder to tell apart. The images need a trailing channel axis before Conv2D will accept them.

import time
import keras

(x_train, y_train), (x_test, y_test) = keras.datasets.fashion_mnist.load_data()
x_train = x_train[..., None].astype("float32") / 255.0    # add channel axis
x_test = x_test[..., None].astype("float32") / 255.0

keras.utils.set_random_seed(0)
model = keras.Sequential([
    keras.Input(shape=(28, 28, 1)),
    keras.layers.Conv2D(32, 3, padding="same", activation="relu"),
    keras.layers.MaxPooling2D(2),
    keras.layers.Conv2D(64, 3, padding="same", activation="relu"),
    keras.layers.MaxPooling2D(2),
    keras.layers.Flatten(),
    keras.layers.Dropout(0.5),
    keras.layers.Dense(10, activation="softmax"),
])
model.compile(optimizer="adam", loss="sparse_categorical_crossentropy",
              metrics=["accuracy"])
print(f"parameters: {model.count_params():,}")
start = time.perf_counter()
model.fit(x_train, y_train, epochs=5, batch_size=128, verbose=2)
print(f"trained in {time.perf_counter() - start:.0f}s")
print(f"test accuracy: {model.evaluate(x_test, y_test, verbose=0)[1]:.4f}")

$ python fashion_cnn.py
parameters: 50,186
Epoch 1/5
469/469 - 10s - 21ms/step - accuracy: 0.7883 - loss: 0.5905
Epoch 2/5
469/469 - 9s - 19ms/step - accuracy: 0.8622 - loss: 0.3871
Epoch 3/5
469/469 - 9s - 19ms/step - accuracy: 0.8783 - loss: 0.3463
Epoch 4/5
469/469 - 9s - 19ms/step - accuracy: 0.8849 - loss: 0.3210
Epoch 5/5
469/469 - 9s - 20ms/step - accuracy: 0.8937 - loss: 0.3010
trained in 46s
test accuracy: 0.8946

Two convolution blocks shrink each 28×28 image to 64 maps of 7×7; Flatten turns those into 3,136 numbers, and the final Dense layer holds 31,370 of the model’s 50,186 parameters. Five epochs on a CPU reach 89.5% test accuracy. Test accuracy is slightly above the last epoch’s training accuracy, partly because Dropout is active while training and switched off for evaluate, and partly because the training figure is averaged over an epoch in which the weights were still improving. The timings are this machine’s; the accuracies repeat exactly on a rerun thanks to set_random_seed.

Data Augmentation Layers

Augmentation layers apply random changes — flips, rotations, zooms — so the model sees a slightly different image every epoch. They only do so in training mode.

import numpy as np
import keras

image = np.arange(9, dtype="float32").reshape(1, 3, 3, 1)
flip = keras.layers.RandomFlip("horizontal", seed=1)

print(image[0, :, :, 0])
print(flip(image, training=False)[0, :, :, 0].numpy())
for _ in range(3):
    print(flip(image, training=True)[0, 0, :, 0].numpy())

$ python augment.py
[[0. 1. 2.]
 [3. 4. 5.]
 [6. 7. 8.]]
[[0. 1. 2.]
 [3. 4. 5.]
 [6. 7. 8.]]
[0. 1. 2.]
[2. 1. 0.]
[0. 1. 2.]

With training=False the image passes through unchanged. With training=True each call flips the image left-to-right with probability one half — the second call reversed the first row. Placed inside a model, fit runs them in training mode and predict and evaluate in inference mode, so augmentation never leaks into predictions.

Loading Images from a Directory

image_dataset_from_directory turns a folder of class subfolders into a labelled tf.data.Dataset. The flower photos dataset has 3,670 images in five folders. Saved as flowers_data.py, this module is imported by the scripts that follow.

from pathlib import Path
import keras

DATA = Path.home() / ".keras/datasets/flower_photos/flower_photos"
if not DATA.exists():
    keras.utils.get_file("flower_photos", extract=True, origin=
        "https://storage.googleapis.com/download.tensorflow.org/example_images/flower_photos.tgz")

train_ds, val_ds = keras.utils.image_dataset_from_directory(
    DATA, image_size=(160, 160), batch_size=32,
    validation_split=0.2, subset="both", seed=42)

if __name__ == "__main__":
    print(train_ds.class_names)
    images, labels = next(iter(train_ds))
    print(images.shape, images.dtype, labels[:8].numpy())
    print(float(images.numpy().min()), float(images.numpy().max()))

$ python flowers_data.py
Found 3670 files belonging to 5 classes.
Using 2936 files for training.
Using 734 files for validation.
['daisy', 'dandelion', 'roses', 'sunflowers', 'tulips']
(32, 160, 160, 3) <dtype: 'float32'> [0 4 0 1 4 3 0 0]
0.0 255.0

Class names come from the folder names in alphabetical order, and each label is the index into that list. subset="both" with a fixed seed returns a training and a validation split that never overlap. Every photo is resized to 160×160, and pixel values stay in [0, 255] — scaling them is the model’s job, as the next example shows.

What a Pretrained Model Already Knows

MobileNetV2 trained on ImageNet ships with Keras; weights="imagenet" downloads its weights (14 MB for the full model) on first use. Ask it about one photo from three of the flower folders.

from pathlib import Path
import numpy as np
import keras
from keras.applications import mobilenet_v2

model = mobilenet_v2.MobileNetV2(weights="imagenet")    # 1000 ImageNet classes
folder = Path.home() / ".keras/datasets/flower_photos/flower_photos"

for name in ["daisy", "sunflowers", "roses"]:
    path = sorted((folder / name).glob("*.jpg"))[0]
    img = keras.utils.load_img(path, target_size=(224, 224))
    x = mobilenet_v2.preprocess_input(np.array(img)[None].astype("float32"))
    top = mobilenet_v2.decode_predictions(model.predict(x, verbose=0), top=2)[0]
    print(f"{name:10s}", ", ".join(f"{label} {p:.2f}" for _, label, p in top))

$ python pretrained.py
daisy      daisy 0.91, bee 0.00
sunflowers daisy 0.79, bee 0.04
roses      daisy 0.25, bee 0.07

The daisy is recognised with 91% confidence — but the sunflower and the rose are called daisies too. ImageNet’s 1,000 classes include daisy and no other flower on this list, so the classifier can only answer with the closest thing it knows. The features underneath still separate flowers well; it is the final 1,000-way layer that is wrong for this task. (load_img needs Pillow: pip install pillow.)

Transfer Learning with a Frozen Base

Keep MobileNetV2’s convolutional base, drop its ImageNet classifier with include_top=False, and train a new five-class head on top.

import time
import keras
from keras.applications import mobilenet_v2
from flowers_data import train_ds, val_ds

keras.utils.set_random_seed(0)
base = mobilenet_v2.MobileNetV2(input_shape=(160, 160, 3), include_top=False,
                                weights="imagenet")
base.trainable = False                            # freeze the ImageNet features
inputs = keras.Input(shape=(160, 160, 3))
x = keras.layers.RandomFlip("horizontal")(inputs)
x = mobilenet_v2.preprocess_input(x)              # [0, 255] -> [-1, 1]
x = base(x, training=False)                       # keep BatchNorm in inference mode
x = keras.layers.GlobalAveragePooling2D()(x)
outputs = keras.layers.Dense(5, activation="softmax")(x)
model = keras.Model(inputs, outputs)

model.compile(optimizer=keras.optimizers.Adam(1e-3),
              loss="sparse_categorical_crossentropy", metrics=["accuracy"])
print(f"trainable: {sum(w.numpy().size for w in model.trainable_weights):,} "
      f"of {model.count_params():,}")
start = time.perf_counter()
model.fit(train_ds, validation_data=val_ds, epochs=3, verbose=2)
print(f"trained in {time.perf_counter() - start:.0f}s")
model.save("flowers.keras")

$ python transfer.py
Found 3670 files belonging to 5 classes.
Using 2936 files for training.
Using 734 files for validation.
trainable: 6,405 of 2,264,389
Epoch 1/3
92/92 - 20s - 218ms/step - accuracy: 0.7299 - loss: 0.7467 - val_accuracy: 0.8338 - val_loss: 0.4673
Epoch 2/3
92/92 - 15s - 162ms/step - accuracy: 0.8784 - loss: 0.3585 - val_accuracy: 0.8638 - val_loss: 0.3733
Epoch 3/3
92/92 - 17s - 189ms/step - accuracy: 0.9121 - loss: 0.2727 - val_accuracy: 0.8569 - val_loss: 0.4077
trained in 52s

Only the new head trains: 1,280 × 5 + 5 = 6,405 weights, against 2.26 million frozen in the base. That is why three epochs over 2,936 photos take about a minute on a CPU and reach 83% validation accuracy after the first. GlobalAveragePooling2D averages each of the base’s 1,280 feature maps down to one number, giving the head a compact feature vector. preprocess_input sits inside the model, so it is saved with it and callers pass raw [0, 255] pixels. Validation accuracy dips slightly in epoch 3 while training accuracy keeps rising, an early sign that the head is fitting the training photos more closely than it generalises.

Fine-Tuning the Top of the Base

With a trained head in place, unfreeze the base’s last 30 layers and continue at a much smaller learning rate, so the pretrained features adapt without being overwritten.

import keras
from flowers_data import train_ds, val_ds

keras.utils.set_random_seed(0)
model = keras.models.load_model("flowers.keras")
base = model.get_layer("mobilenetv2_1.00_160")
print(f"base has {len(base.layers)} layers")

base.trainable = True
for layer in base.layers[:-30]:                   # unfreeze only the last 30
    layer.trainable = False

model.compile(optimizer=keras.optimizers.Adam(1e-5),   # 100x smaller step
              loss="sparse_categorical_crossentropy", metrics=["accuracy"])
print(f"trainable: {sum(w.numpy().size for w in model.trainable_weights):,}")
print(f"before: accuracy {model.evaluate(train_ds, verbose=0)[1]:.4f} - "
      f"val_accuracy {model.evaluate(val_ds, verbose=0)[1]:.4f}")
model.fit(train_ds, validation_data=val_ds, epochs=3, verbose=2)

$ python finetune.py
Found 3670 files belonging to 5 classes.
Using 2936 files for training.
Using 734 files for validation.
base has 154 layers
trainable: 1,532,805
before: accuracy 0.9251 - val_accuracy 0.8569
Epoch 1/3
92/92 - 29s - 311ms/step - accuracy: 0.8580 - loss: 0.4184 - val_accuracy: 0.8638 - val_loss: 0.3900
Epoch 2/3
92/92 - 21s - 228ms/step - accuracy: 0.8965 - loss: 0.3114 - val_accuracy: 0.8678 - val_loss: 0.3636
Epoch 3/3
92/92 - 23s - 254ms/step - accuracy: 0.9159 - loss: 0.2651 - val_accuracy: 0.8706 - val_loss: 0.3444

Any change to trainable needs a fresh compile to take effect. Unfreezing 30 layers raises the trainable count from 6,405 to 1.5 million, and validation accuracy climbs from 85.7% to 87.1% over three epochs. Training accuracy first drops — from the 92.5% the loaded model scores on the same photos to 85.8% averaged over the first fine-tuning epoch — because the first updates to the newly unfrozen weights disturb features the head was fitted to; it recovers by epoch 3. (A control run of one epoch with the learning rate set to 0 stayed at 92.4%, so the drop comes from the updates, not from RandomFlip or epoch averaging.) The BatchNormalization layers among the unfrozen 30 stay in inference mode because the base was called with training=False, so their ImageNet statistics are not overwritten by small flower batches. The gain is modest here; fine-tuning pays off most when your images differ more from ImageNet than flower photos do.

Practice

import numpy as np
import keras

conv = keras.layers.Conv2D(16, 5, strides=2, padding="same")
x = np.zeros((1, 32, 32, 3), dtype="float32")

print(conv(x).shape)
print(conv.count_params(), "=", 5 * 5 * 3 * 16 + 16)

$ python p1_shape.py
(1, 16, 16, 16)
1216 = 1216

import numpy as np
import keras

image = np.zeros((1, 6, 5, 1), dtype="float32")
image[:, 3:, :, :] = 1.0                 # dark top, bright bottom

conv = keras.layers.Conv2D(1, 3, use_bias=False)
conv.build(image.shape)
kernel = np.array([[-1, -1, -1],
                   [ 0,  0,  0],
                   [ 1,  1,  1]], dtype="float32")   # bottom minus top
conv.set_weights([kernel.reshape(3, 3, 1, 1)])
print(conv(image).numpy()[0, :, :, 0])

$ python p2_horizontal.py
[[0. 0. 0.]
 [3. 3. 3.]
 [3. 3. 3.]
 [0. 0. 0.]]

import keras

def build(head):
    return keras.Sequential([
        keras.Input(shape=(28, 28, 1)),
        keras.layers.Conv2D(32, 3, padding="same", activation="relu"),
        keras.layers.MaxPooling2D(2),
        keras.layers.Conv2D(64, 3, padding="same", activation="relu"),
        keras.layers.MaxPooling2D(2),
        head,
        keras.layers.Dense(10, activation="softmax"),
    ])

print(f"Flatten:                {build(keras.layers.Flatten()).count_params():,}")
print(f"GlobalAveragePooling2D: {build(keras.layers.GlobalAveragePooling2D()).count_params():,}")

$ python p3_gap.py
Flatten:                50,186
GlobalAveragePooling2D: 19,466

import numpy as np
import keras
from flowers_data import val_ds

model = keras.models.load_model("flowers.keras")
path = next(p for p in val_ds.file_paths if "/roses/" in p)   # held-out photo
img = keras.utils.load_img(path, target_size=(160, 160))
probs = model.predict(np.array(img)[None].astype("float32"), verbose=0)[0]

for name, p in zip(val_ds.class_names, probs):
    print(f"{name:10s} {p:.3f}")

$ python p4_predict.py
Found 3670 files belonging to 5 classes.
Using 2936 files for training.
Using 734 files for validation.
daisy      0.012
dandelion  0.000
roses      0.988
sunflowers 0.000
tulips     0.000