TensorFlow · lesson 2 of 7
Also part of: Machine Learning
Introduction to TensorFlow and Keras
Learn TensorFlow tensors, GradientTape, and Keras 3 through runnable examples: train an MNIST classifier, a functional model, and a custom training loop.
By SkillByExample · Updated
What Is TensorFlow?
TensorFlow is Google’s open-source machine learning framework. It runs eagerly by default — each line executes immediately, like NumPy — and can compile Python functions into graphs with tf.function for speed and deployment. Keras is its high-level model API; since TensorFlow 2.16, tf.keras is Keras 3.
Installation
pip install tensorflow
The examples below ran with TensorFlow 2.21 and Keras 3.15 on a CPU-only machine. On import, TensorFlow writes several log lines to stderr — that no CUDA driver was found, that oneDNN is enabled — and then runs on the CPU. Those lines are left out of the outputs shown here; everything else is exactly what each program printed.
Tensors
A tensor is an immutable n-dimensional array with a shape and a dtype.
import tensorflow as tf
x = tf.constant([[1.0, 2.0], [3.0, 4.0]])
print(x.shape, x.dtype)
print(x @ x) # matrix multiply
print(tf.reduce_sum(x, axis=0).numpy())
ints = tf.constant([1, 2, 3])
print(ints.dtype)
try:
x[0, 0] = 9.0
except TypeError as e:
print("TypeError:", e)
v = tf.Variable(x)
v[0, 0].assign(9.0)
print(v.numpy())
$ python tensors.py
(2, 2) <dtype: 'float32'>
tf.Tensor(
[[ 7. 10.]
[15. 22.]], shape=(2, 2), dtype=float32)
[4. 6.]
<dtype: 'int32'>
TypeError: 'tensorflow.python.framework.ops.EagerTensor' object does not support item assignment
[[9. 2.]
[3. 4.]]
Python floats become float32 and Python ints int32 — not NumPy’s 64-bit defaults. @ is matrix multiplication, so x @ x is [[1·1+2·3, 1·2+2·4], …], not element-wise squaring. .numpy() converts to a NumPy array. Tensors cannot be changed in place; a tf.Variable can, through .assign, which is why model weights are Variables.
Gradients with GradientTape
tf.GradientTape records operations as they run, then computes derivatives backwards through them.
import tensorflow as tf
x = tf.Variable(3.0)
with tf.GradientTape() as tape:
y = x ** 2 + 2 * x # dy/dx = 2x + 2
print(tape.gradient(y, x).numpy())
c = tf.constant(3.0)
with tf.GradientTape() as tape:
y = c ** 2
print(tape.gradient(y, c))
with tf.GradientTape() as tape:
tape.watch(c)
y = c ** 2
print(tape.gradient(y, c).numpy())
$ python tape.py
8.0
None
6.0
At x = 3, 2x + 2 = 8. The second gradient is None, not an error: the tape only tracks trainable Variables automatically, and c is a constant. tape.watch(c) opts it in, and the derivative 2c = 6 appears. A None gradient in real training usually means a weight was not a Variable or was not used to compute the loss.
Your First Keras Model
MNIST is 70,000 28×28 grayscale digits; load_data downloads it on first use. Flatten each image to 784 numbers, scale to [0, 1], and train a one-hidden-layer network.
import keras
(x_train, y_train), (x_test, y_test) = keras.datasets.mnist.load_data()
x_train = x_train.reshape(-1, 784).astype("float32") / 255.0
x_test = x_test.reshape(-1, 784).astype("float32") / 255.0
print(x_train.shape, y_train[:5])
keras.utils.set_random_seed(0)
model = keras.Sequential([
keras.Input(shape=(784,)),
keras.layers.Dense(128, activation="relu"),
keras.layers.Dense(10, activation="softmax"),
])
model.compile(optimizer="adam",
loss="sparse_categorical_crossentropy",
metrics=["accuracy"])
model.fit(x_train, y_train, epochs=3, batch_size=128,
validation_split=0.1, verbose=2)
loss, acc = model.evaluate(x_test, y_test, verbose=0)
print(f"test accuracy: {acc:.4f}")
print(model.predict(x_test[:1], verbose=0).round(3))
$ python mnist_mlp.py
(60000, 784) [5 0 4 1 9]
Epoch 1/3
422/422 - 2s - 6ms/step - accuracy: 0.8991 - loss: 0.3718 - val_accuracy: 0.9542 - val_loss: 0.1702
Epoch 2/3
422/422 - 1s - 3ms/step - accuracy: 0.9495 - loss: 0.1769 - val_accuracy: 0.9655 - val_loss: 0.1237
Epoch 3/3
422/422 - 1s - 3ms/step - accuracy: 0.9641 - loss: 0.1266 - val_accuracy: 0.9702 - val_loss: 0.1048
test accuracy: 0.9657
[[0. 0. 0. 0.002 0. 0. 0. 0.997 0. 0. ]]
validation_split=0.1 holds back the last 6,000 training images, leaving 54,000 — at 128 per batch that is the 422 steps per epoch. sparse_categorical_crossentropy takes integer labels like 5 directly, with no one-hot encoding. Training accuracy is averaged over the epoch while the weights are still improving, which is why it trails validation accuracy here. The prediction is a softmax over the 10 digits: 99.7% on 7, the correct label for the first test image. Timings depend on the machine.
Reading model.summary()
summary() prints each layer’s output shape and parameter count — the quickest check that a model is wired the way you intended.
import keras
model = keras.Sequential([
keras.Input(shape=(784,)),
keras.layers.Dense(128, activation="relu"),
keras.layers.Dense(10, activation="softmax"),
])
model.summary()
$ python summary.py
Model: "sequential"
┏━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━┳━━━━━━━━━━━━━━━━━━━━━━━━┳━━━━━━━━━━━━━━━┓
┃ Layer (type) ┃ Output Shape ┃ Param # ┃
┡━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━╇━━━━━━━━━━━━━━━━━━━━━━━━╇━━━━━━━━━━━━━━━┩
│ dense (Dense) │ (None, 128) │ 100,480 │
├─────────────────────────────────┼────────────────────────┼───────────────┤
│ dense_1 (Dense) │ (None, 10) │ 1,290 │
└─────────────────────────────────┴────────────────────────┴───────────────┘
Total params: 101,770 (397.54 KB)
Trainable params: 101,770 (397.54 KB)
Non-trainable params: 0 (0.00 B)
None is the batch dimension, left open so any batch size works. A Dense layer has one weight per input-output pair plus one bias per output: 784 × 128 + 128 = 100,480 and 128 × 10 + 10 = 1,290. At 4 bytes per float32, 101,770 parameters is 397.54 KB. keras.Input is what lets Keras build the layers — and count parameters — before any data arrives.
The Functional API
When a model has more than one input, wire the layers by calling them on tensors and hand the inputs and outputs to keras.Model.
import numpy as np
import keras
text = keras.Input(shape=(100,), name="text")
meta = keras.Input(shape=(4,), name="meta")
x = keras.layers.Dense(16, activation="relu")(text)
merged = keras.layers.Concatenate()([x, meta])
out = keras.layers.Dense(1, activation="sigmoid", name="score")(merged)
model = keras.Model(inputs=[text, meta], outputs=out)
for layer in model.layers:
print(f"{layer.name:12s} {str(layer.output.shape):10s} {layer.count_params()}")
rng = np.random.default_rng(0)
batch = {"text": rng.random((3, 100)), "meta": rng.random((3, 4))}
print(model.predict(batch, verbose=0).shape)
$ python functional.py
text (None, 100) 0
dense (None, 16) 1616
meta (None, 4) 0
concatenate (None, 20) 0
score (None, 1) 21
(3, 1)
The text branch compresses 100 features to 16 (100 × 16 + 16 = 1,616 parameters); Concatenate joins those with the 4 metadata features into 20, and the output layer maps 20 to one score (20 + 1 = 21). Because the inputs are named, predict accepts a dict keyed by those names, so the order of the arrays cannot be mixed up. Inputs and Concatenate have no parameters — they only route tensors.
A Custom Training Loop
model.fit is a loop over batches that computes a loss, takes its gradient, and applies an optimizer step. Writing that loop by hand shows each piece — here fitting y = 2a − b + 0.5c + 3.
import numpy as np
import tensorflow as tf
import keras
rng = np.random.default_rng(0)
X = rng.normal(size=(1000, 3)).astype("float32")
y = X @ np.array([[2.0], [-1.0], [0.5]], dtype="float32") + 3.0
model = keras.Sequential([keras.Input(shape=(3,)), keras.layers.Dense(1)])
optimizer = keras.optimizers.SGD(learning_rate=0.1)
loss_fn = keras.losses.MeanSquaredError()
ds = tf.data.Dataset.from_tensor_slices((X, y)).shuffle(1000, seed=0).batch(100)
@tf.function
def train_step(xb, yb):
with tf.GradientTape() as tape:
loss = loss_fn(yb, model(xb, training=True))
grads = tape.gradient(loss, model.trainable_variables)
optimizer.apply_gradients(zip(grads, model.trainable_variables))
return loss
for epoch in range(3):
losses = [train_step(xb, yb).numpy() for xb, yb in ds]
print(f"epoch {epoch}: loss {np.mean(losses):.4f}")
print(model.layers[0].kernel.numpy().ravel().round(3), model.layers[0].bias.numpy().round(3))
$ python loop.py
epoch 0: loss 5.6474
epoch 1: loss 0.1092
epoch 2: loss 0.0021
[ 1.991 -0.999 0.498] [2.991]
After three passes over the data the single Dense layer has recovered the true coefficients 2, -1, 0.5 and the intercept 3 to within 0.01. tf.data.Dataset handles shuffling and batching; @tf.function traces train_step once into a graph and reuses it for every batch. model.trainable_variables is the list the tape differentiates against — the kernel and the bias.
Saving and Loading
import os
import numpy as np
import keras
keras.utils.set_random_seed(0)
model = keras.Sequential([keras.Input(shape=(4,)),
keras.layers.Dense(8, activation="relu"),
keras.layers.Dense(1)])
model.compile(optimizer="adam", loss="mse")
x = np.ones((2, 4), dtype="float32")
model.save("model.keras")
restored = keras.models.load_model("model.keras")
print(np.allclose(model.predict(x, verbose=0), restored.predict(x, verbose=0)))
print(os.path.getsize("model.keras"), "bytes")
try:
model.save("saved_model_dir")
except ValueError as e:
print("ValueError:", e)
$ python save.py
True
18540 bytes
ValueError: Invalid filepath extension for saving. Please add either a `.keras` extension for the native Keras format (recommended) or a `.h5` extension. Use `model.export(filepath)` if you want to export a SavedModel for use with TFLite/TFServing/etc. Received: filepath=saved_model_dir.
A .keras file is a zip holding the architecture, the weights, and the compile settings, so the restored model gives identical predictions and can resume training. Older tutorials call model.save("dir") to get a TensorFlow SavedModel; in Keras 3 that raises, and the error names the replacement — model.export, covered in TensorFlow Model Deployment.
Practice
1. Use GradientTape to find the derivative of x ** 3 at x = 2. Predict the answer first.
import tensorflow as tf
x = tf.Variable(2.0)
with tf.GradientTape() as tape:
y = x ** 3 # dy/dx = 3x^2
print(tape.gradient(y, x).numpy())
$ python p1_cube.py
12.0
3x² = 3 × 4 = 12. The tape computes this numerically by walking back through the pow operation it recorded; it never forms the symbolic expression.
2. Add a Dropout(0.2) and a second hidden layer of 64 units to the MNIST model, and count the parameters layer by layer.
import keras
model = keras.Sequential([
keras.Input(shape=(784,)),
keras.layers.Dense(128, activation="relu"),
keras.layers.Dropout(0.2),
keras.layers.Dense(64, activation="relu"),
keras.layers.Dense(10, activation="softmax"),
])
for layer in model.layers:
print(f"{layer.name:10s} {layer.count_params():>7,}")
print(f"total {model.count_params():>7,}")
$ python p2_params.py
dense 100,480
dropout 0
dense_1 8,256
dense_2 650
total 109,386
Dropout has no weights — it only zeroes a random 20% of activations during training. The new layer adds 128 × 64 + 64 = 8,256, and the output layer shrinks to 64 × 10 + 10 = 650 because it now reads 64 inputs instead of 128.
3. In loop.py, lower the learning rate to 0.01. What happens to the loss after three epochs?
import numpy as np
import tensorflow as tf
import keras
rng = np.random.default_rng(0)
X = rng.normal(size=(1000, 3)).astype("float32")
y = X @ np.array([[2.0], [-1.0], [0.5]], dtype="float32") + 3.0
model = keras.Sequential([keras.Input(shape=(3,)), keras.layers.Dense(1)])
optimizer = keras.optimizers.SGD(learning_rate=0.01)
loss_fn = keras.losses.MeanSquaredError()
ds = tf.data.Dataset.from_tensor_slices((X, y)).shuffle(1000, seed=0).batch(100)
@tf.function
def train_step(xb, yb):
with tf.GradientTape() as tape:
loss = loss_fn(yb, model(xb, training=True))
grads = tape.gradient(loss, model.trainable_variables)
optimizer.apply_gradients(zip(grads, model.trainable_variables))
return loss
for epoch in range(3):
losses = [train_step(xb, yb).numpy() for xb, yb in ds]
print(f"epoch {epoch}: loss {np.mean(losses):.4f}")
print(model.layers[0].kernel.numpy().ravel().round(3), model.layers[0].bias.numpy().round(3))
$ python p3_lr.py
epoch 0: loss 16.2188
epoch 1: loss 11.2959
epoch 2: loss 7.8700
[0.378 0.206 0.621] [1.248]
Each step now moves the weights a tenth as far, so after the same 30 steps the model is still far from the answer — the bias is 1.25 instead of 3. The loss is still falling steadily; it simply needs many more epochs.
4. Turn the MNIST model's softmax outputs into predicted digits and find the first test image it gets wrong.
import numpy as np
import keras
(x_train, y_train), (x_test, y_test) = keras.datasets.mnist.load_data()
x_train = x_train.reshape(-1, 784).astype("float32") / 255.0
x_test = x_test.reshape(-1, 784).astype("float32") / 255.0
keras.utils.set_random_seed(0)
model = keras.Sequential([
keras.Input(shape=(784,)),
keras.layers.Dense(128, activation="relu"),
keras.layers.Dense(10, activation="softmax"),
])
model.compile(optimizer="adam", loss="sparse_categorical_crossentropy",
metrics=["accuracy"])
model.fit(x_train, y_train, epochs=3, batch_size=128, verbose=0)
pred = np.argmax(model.predict(x_test, verbose=0), axis=1)
print("predicted:", pred[:10])
print("actual: ", y_test[:10])
wrong = np.flatnonzero(pred != y_test)
print(f"{len(wrong)} wrong of {len(y_test)}; first: index {wrong[0]}, "
f"predicted {pred[wrong[0]]}, actual {y_test[wrong[0]]}")
$ python p4_argmax.py
predicted: [7 2 1 0 4 1 4 9 6 9]
actual: [7 2 1 0 4 1 4 9 5 9]
320 wrong of 10000; first: index 8, predicted 6, actual 5
argmax over axis 1 picks the most probable digit per row. 320 errors is 96.8% accuracy — slightly better than the 96.57% above because this run trains on all 60,000 images instead of holding 6,000 back for validation.