Skip to main content
TensorFlow · lesson 2 of 7
TensorFlow beginner 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.

Frequently Asked Questions

What is the difference between TensorFlow and Keras?
TensorFlow is the numerical engine: tensors, automatic differentiation with GradientTape, graph compilation with tf.function, and deployment tools such as SavedModel and TF Serving. Keras is the high-level API for defining, training, and saving models. Since TensorFlow 2.16, tf.keras is Keras 3, which can also run on JAX or PyTorch.
When should I use the Sequential API vs the Functional API?
Use Sequential for a plain stack of layers where each layer has one input and one output. Use the Functional API when a model has several inputs or outputs, shared layers, or skip connections, because you wire the layers together yourself. A Functional model can do everything a Sequential one can.
Why does TensorFlow print CUDA errors when I have no GPU?
The standard tensorflow pip package includes GPU support and tries to load the CUDA driver at import time. On a machine without an NVIDIA GPU it logs that the driver was not found and falls back to the CPU. The messages go to stderr and are harmless; your code runs normally on the CPU.
Why does model.save('my_model') raise a ValueError in Keras 3?
Keras 3 requires a file extension so it knows which format to write: use model.save('my_model.keras') for the native Keras format. To produce a TensorFlow SavedModel directory for TF Serving or TFLite conversion, call model.export('my_model') instead.
When do I need a custom training loop instead of model.fit()?
model.fit() handles batching, metrics, callbacks, and validation, and covers most training needs. Write a loop with tf.GradientTape when the update rule itself is unusual, such as a GAN with two optimizers, gradient penalties, or per-step logic that callbacks cannot express. Wrap the step in tf.function so it runs as a compiled graph.