Week 1 spent a day writing a training loop by hand and week 2 spent a day making it respectable. Keras replaces the whole of that with two method calls, and it is worth being precise about what you gain and what you give up.
import os
os.environ['KERAS_BACKEND'] = 'tensorflow'
import numpy as np
import keras
from keras import layers
(xtr, ytr), (xte, yte) = keras.datasets.mnist.load_data()
xtr = (xtr[:12000].astype('float32') / 255.0)[..., None]
ytr = ytr[:12000]
xte = (xte[:2000].astype('float32') / 255.0)[..., None]
yte = yte[:2000]
keras.utils.set_random_seed(0)
def build_keras():
return keras.Sequential([
layers.Input((28, 28, 1)),
layers.Conv2D(32, 3, padding='same', activation='relu'),
layers.MaxPooling2D(),
layers.Conv2D(64, 3, padding='same', activation='relu'),
layers.MaxPooling2D(),
layers.Flatten(),
layers.Dense(64, activation='relu'),
layers.Dense(10)])
import time
keras.utils.set_random_seed(0)
model = build_keras()
model.compile(optimizer=keras.optimizers.Adam(1e-3),
loss=keras.losses.SparseCategoricalCrossentropy(
from_logits=True),
metrics=['accuracy'])
start = time.perf_counter()
hist = model.fit(xtr, ytr, epochs=5, batch_size=128, verbose=2,
validation_data=(xte, yte))
print()
print('keras took %.1f seconds' % (time.perf_counter() - start))
print('final validation accuracy %.4f'
% hist.history['val_accuracy'][-1])
Epoch 1/5
94/94 - 5s - 48ms/step - accuracy: 0.8140 - loss: 0.6536 - val_accuracy: 0.9250 - val_loss: 0.2761
Epoch 2/5
94/94 - 4s - 39ms/step - accuracy: 0.9540 - loss: 0.1589 - val_accuracy: 0.9555 - val_loss: 0.1538
Epoch 3/5
94/94 - 4s - 40ms/step - accuracy: 0.9732 - loss: 0.0934 - val_accuracy: 0.9630 - val_loss: 0.1117
Epoch 4/5
94/94 - 4s - 40ms/step - accuracy: 0.9812 - loss: 0.0661 - val_accuracy: 0.9680 - val_loss: 0.0929
Epoch 5/5
94/94 - 4s - 40ms/step - accuracy: 0.9866 - loss: 0.0525 - val_accuracy: 0.9705 - val_loss: 0.0823
keras took 19.6 seconds
final validation accuracy 0.9705
Everything week 2 insisted on is in there. The validation split is evaluated every epoch, the metrics are tracked separately for training and validation, the model is switched between training and evaluation mode around the right calls, and the history is kept. None of it was written.
from_logits=True is not optional
The model above ends in a plain Dense(10) with no activation, so its outputs are logits. Week 2 made the case for that, and the loss has to be told. Leave the flag off and Keras takes the logits for probabilities, applies a log to numbers that may be negative, and trains something that is not what you meant. It does not raise anything.
What compile is actually doing
- Attaching the optimiser, so
fit knows what to step.
- Attaching the loss, and building the metric objects that accumulate across a whole epoch rather than reporting the last batch.
- On the TensorFlow backend, tracing the training step into a graph, which is why the first epoch above is slower than the rest.
That last point explains a timing result people misread constantly. The first epoch pays for compilation, so a benchmark of one epoch makes Keras look worse than it is, and a benchmark that ignores it makes it look better.
Writing your own layer
A custom layer is the same shape of object as an nn.Module: declare the weights, say what the forward pass does. Keras splits it into build, which runs once when the input shape is known, and call, which is the forward pass. That split is what lets you write a layer without knowing its input width, which PyTorch makes you supply.
import os
os.environ['KERAS_BACKEND'] = 'tensorflow'
import numpy as np
import keras
from keras import layers
(xtr, ytr), (xte, yte) = keras.datasets.mnist.load_data()
xtr = (xtr[:12000].astype('float32') / 255.0)[..., None]
ytr = ytr[:12000]
xte = (xte[:2000].astype('float32') / 255.0)[..., None]
yte = yte[:2000]
keras.utils.set_random_seed(0)
class Scale(keras.layers.Layer):
"""One learned multiplier per feature."""
def build(self, input_shape):
self.w = self.add_weight(shape=(input_shape[-1],),
initializer='ones', trainable=True)
def call(self, x):
return x * self.w
layer = Scale()
out = layer(np.ones((2, 5), dtype='float32'))
print('output %s' % (out.shape,))
print('weights created on first call: %s'
% [tuple(w.shape) for w in layer.weights])
print('trainable: %d' % len(layer.trainable_weights))
output (2, 5)
weights created on first call: [(5,)]
trainable: 1
Where the abstraction stops paying
Everything above holds as long as your training step is the ordinary one. When it is not, you override train_step, and the code inside it is written against a particular backend rather than against Keras.
import os
os.environ['KERAS_BACKEND'] = 'tensorflow'
import numpy as np
import keras
from keras import layers
(xtr, ytr), (xte, yte) = keras.datasets.mnist.load_data()
xtr = (xtr[:12000].astype('float32') / 255.0)[..., None]
ytr = ytr[:12000]
xte = (xte[:2000].astype('float32') / 255.0)[..., None]
yte = yte[:2000]
keras.utils.set_random_seed(0)
import tensorflow as tf
class Noisy(keras.Model):
"""An ordinary model that adds noise to its inputs while training."""
def __init__(self, inner):
super().__init__()
self.inner = inner
def call(self, x, training=False):
return self.inner(x)
def train_step(self, data):
x, y = data
x = x + tf.random.normal(tf.shape(x), stddev=0.1)
with tf.GradientTape() as tape:
loss = self.compute_loss(y=y, y_pred=self(x, training=True))
grads = tape.gradient(loss, self.trainable_variables)
self.optimizer.apply_gradients(zip(grads,
self.trainable_variables))
return {'loss': loss}
keras.utils.set_random_seed(0)
model = Noisy(keras.Sequential([keras.layers.Input((28, 28, 1)),
keras.layers.Flatten(),
keras.layers.Dense(10)]))
model.compile(optimizer='adam',
loss=keras.losses.SparseCategoricalCrossentropy(
from_logits=True))
hist = model.fit(xtr[:2000], ytr[:2000], epochs=2, batch_size=128,
verbose=2)
print()
print('the loop ran, and every line inside train_step was TensorFlow')
Epoch 1/2
16/16 - 0s - 17ms/step - loss: 0.0000e+00
Epoch 2/2
16/16 - 0s - 4ms/step - loss: 0.0000e+00
the loop ran, and every line inside train_step was TensorFlow
A custom train_step is not portable across backends
tf.GradientTape is TensorFlow. On the PyTorch backend the same method is written with loss.backward() and an optimiser step, and on JAX it is a pure function returning gradients. So the moment you need a training step Keras does not already provide, you have picked a backend and the portability that was the reason for Keras 3 is gone.
This is the honest boundary of the abstraction, and it is worth knowing before you build on it rather than after.