Example 4: TensorFlow + CIFAR-10 + checkpoint

This example shows how to checkpoint and resume a GPU deep learning job. It uses TensorFlow to train a small convolutional neural network on the CIFAR-10 dataset. The model uses TensorFlow’s tf.keras layers and metrics, while checkpointing is handled with TensorFlow tf.train checkpoints. The job saves a checkpoint after each epoch, so if the job stops before training finishes, you can submit it again and continue from the latest checkpoint.

The example runs inside a public TensorFlow GPU container from Docker Hub: tensorflow/tensorflow:latest.

1. Connect to the CoSTAR cluster

Open a terminal and connect to the CoSTAR login node:

ssh ab1234@costar-login01

Replace ab1234 with your University username.

2. Create the submission script

Download the example submission script, or create a file called example-04-checkpoint.sh in your working directory and copy in the script below.

example-04-checkpoint.sh (click to collapse / view)
 1#!/bin/bash
 2
 3#SBATCH --job-name=tf-checkpoint
 4#SBATCH --partition=main
 5#SBATCH --gpus=1
 6#SBATCH --nodes=1
 7#SBATCH --ntasks=1
 8#SBATCH --cpus-per-task=1
 9#SBATCH --mem=20G
10#SBATCH --time=00:45:00
11#SBATCH --output=slurm.%N.%j.out
12#SBATCH --error=slurm.%N.%j.err
13
14# Training variables
15script=$PWD/example-04-checkpoint.py
16ckpt_dir=$PWD/models_tf
17
18# Hyperparameters
19batch_size=32
20epochs=10
21lr=0.01
22momentum=0.5
23
24apptainer exec docker://tensorflow/tensorflow:latest \
25    python "$script" \
26    --ckpt-dir "$ckpt_dir" \
27    --batch-size "$batch_size" \
28    --epochs "$epochs" \
29    --lr "$lr" \
30    --momentum "$momentum" \
31    --resume-training

The docker://tensorflow/tensorflow:latest image reference tells Apptainer to pull the public TensorFlow GPU container from Docker Hub and run the Python script inside that container.

Note

The first run may take longer while Apptainer pulls and converts the Docker image. Later runs are usually faster because Apptainer can reuse the cached image.

3. Create the Python script

Download the Python script, or create a file called example-04-checkpoint.py in your working directory and copy in the script below.

example-04-checkpoint.py (click to collapse / view)
  1from __future__ import annotations
  2import argparse
  3from pathlib import Path
  4import tensorflow as tf
  5
  6
  7def parse_args():
  8    parser = argparse.ArgumentParser(description="TensorFlow CIFAR-10 CNN with checkpointing")
  9    parser.add_argument('--batch-size', type=int, default=64, metavar='N',
 10                        help='input batch size for training (default: 64)')
 11    parser.add_argument('--test-batch-size', type=int, default=1000, metavar='N',
 12                        help='input batch size for testing (default: 1000)')
 13    parser.add_argument('--epochs', type=int, default=10, metavar='N',
 14                        help='number of epochs to train (default: 10)')
 15    parser.add_argument('--lr', type=float, default=0.01, metavar='LR',
 16                        help='learning rate (default: 0.01)')
 17    parser.add_argument('--momentum', type=float, default=0.5, metavar='M',
 18                        help='SGD momentum (default: 0.5)')
 19    parser.add_argument('--no-cuda', action='store_true', default=False,
 20                        help='disables CUDA training')
 21    parser.add_argument('--seed', type=int, default=1, metavar='S',
 22                        help='random seed (default: 1)')
 23    parser.add_argument('--log-interval', type=int, default=10, metavar='N',
 24                        help='how many batches to wait before logging training status')
 25    parser.add_argument('--ckpt-dir', required=True, help='path to save and load checkpoints')
 26    parser.add_argument('--resume-training', action='store_true', help='resume training from latest checkpoint')
 27    return parser.parse_args()
 28
 29
 30def disable_gpu_if_requested(no_cuda: bool) -> None:
 31    if not no_cuda:
 32        return
 33    try:
 34        tf.config.set_visible_devices([], 'GPU')
 35        print('GPU devices disabled via --no-cuda; using CPU.')
 36    except Exception as exc:  # pragma: no cover - defensive
 37        print(f'Could not disable GPU devices: {exc}')
 38
 39
 40def build_datasets(batch_size: int, test_batch_size: int, seed: int):
 41    (x_train, y_train), (x_test, y_test) = tf.keras.datasets.cifar10.load_data()
 42    x_train = x_train.astype('float32') / 255.0
 43    x_test = x_test.astype('float32') / 255.0
 44
 45    y_train = y_train.reshape(-1)
 46    y_test = y_test.reshape(-1)
 47
 48    train_ds = (tf.data.Dataset.from_tensor_slices((x_train, y_train))
 49                .shuffle(buffer_size=10000, seed=seed)
 50                .batch(batch_size)
 51                .prefetch(tf.data.AUTOTUNE))
 52
 53    test_ds = (tf.data.Dataset.from_tensor_slices((x_test, y_test))
 54               .batch(test_batch_size)
 55               .prefetch(tf.data.AUTOTUNE))
 56    return train_ds, test_ds
 57
 58
 59def create_model():
 60    return tf.keras.Sequential([
 61        tf.keras.layers.Input(shape=(32, 32, 3)),
 62        tf.keras.layers.Conv2D(32, 3, activation='relu'),
 63        tf.keras.layers.MaxPooling2D(pool_size=(2, 2)),
 64        tf.keras.layers.Conv2D(64, 3, activation='relu'),
 65        tf.keras.layers.Dropout(0.3),
 66        tf.keras.layers.MaxPooling2D(pool_size=(2, 2)),
 67        tf.keras.layers.Flatten(),
 68        tf.keras.layers.Dense(50, activation='relu'),
 69        tf.keras.layers.Dropout(0.5),
 70        tf.keras.layers.Dense(10)
 71    ])
 72
 73
 74def main():
 75    args = parse_args()
 76    tf.random.set_seed(args.seed)
 77    disable_gpu_if_requested(args.no_cuda)
 78
 79    ckpt_dir = Path(args.ckpt_dir)
 80    ckpt_dir.mkdir(parents=True, exist_ok=True)
 81
 82    train_ds, test_ds = build_datasets(args.batch_size, args.test_batch_size, args.seed)
 83
 84    model = create_model()
 85    loss_fn = tf.keras.losses.SparseCategoricalCrossentropy(from_logits=True)
 86    optimizer = tf.keras.optimizers.SGD(learning_rate=args.lr, momentum=args.momentum)
 87
 88    ckpt = tf.train.Checkpoint(epoch=tf.Variable(0, dtype=tf.int64),
 89                               optimizer=optimizer,
 90                               model=model)
 91    manager = tf.train.CheckpointManager(ckpt, ckpt_dir, max_to_keep=5)
 92
 93    start_epoch = 0
 94    if args.resume_training and manager.latest_checkpoint:
 95        ckpt.restore(manager.latest_checkpoint)
 96        start_epoch = int(ckpt.epoch.numpy())
 97        print(f'Resuming from checkpoint: {manager.latest_checkpoint} (start at epoch {start_epoch + 1})')
 98    elif args.resume_training:
 99        print(f'No checkpoints found in {ckpt_dir}. Starting from scratch...')
100
101    @tf.function
102    def train_step(images, labels):
103        with tf.GradientTape() as tape:
104            logits = model(images, training=True)
105            loss_value = loss_fn(labels, logits)
106        grads = tape.gradient(loss_value, model.trainable_variables)
107        optimizer.apply_gradients(zip(grads, model.trainable_variables))
108        return loss_value, logits
109
110    @tf.function
111    def test_step(images, labels):
112        logits = model(images, training=False)
113        t_loss = loss_fn(labels, logits)
114        return t_loss, logits
115
116    for epoch in range(start_epoch, args.epochs):
117        train_loss = tf.keras.metrics.Mean()
118        train_accuracy = tf.keras.metrics.SparseCategoricalAccuracy()
119        test_loss = tf.keras.metrics.Mean()
120        test_accuracy = tf.keras.metrics.SparseCategoricalAccuracy()
121
122        for batch_idx, (images, labels) in enumerate(train_ds):
123            loss_value, logits = train_step(images, labels)
124            train_loss.update_state(loss_value)
125            train_accuracy.update_state(labels, logits)
126            if batch_idx % args.log_interval == 0:
127                print(f'Train Epoch: {epoch + 1} [{batch_idx * len(images):>5d}]\tLoss: {loss_value.numpy():.6f}')
128
129        for images, labels in test_ds:
130            t_loss, logits = test_step(images, labels)
131            test_loss.update_state(t_loss)
132            test_accuracy.update_state(labels, logits)
133
134        ckpt.epoch.assign(epoch + 1)
135        ckpt_path = manager.save()
136        print(f'Saved checkpoint to {ckpt_path}')
137
138        print(f'\nEpoch {epoch + 1}/{args.epochs}: ')
139        print(f'  Train   - loss: {train_loss.result():.4f}, accuracy: {train_accuracy.result() * 100:.2f}%')
140        print(f'  Test    - loss: {test_loss.result():.4f}, accuracy: {test_accuracy.result() * 100:.2f}%\n')
141
142
143if __name__ == '__main__':
144    main()

4. Submit your job

Submit the script to the Slurm scheduler:

sbatch example-04-checkpoint.sh

You can check the status of your job by running:

squeue --me

5. Check the output and checkpoints

When the job finishes, Slurm writes the standard output and error streams to files in the submission directory:

  • slurm.<node>.<job_id>.out contains TensorFlow GPU detection messages, training progress, and checkpoint messages.

  • slurm.<node>.<job_id>.err contains error messages, if any were produced.

The submission script sets the checkpoint directory with ckpt_dir and passes it to the Python script using --ckpt-dir. In this example, checkpoints are written to models_tf/ in your working directory:

models_tf/ckpt-1
models_tf/ckpt-2
models_tf/ckpt-3

These files are written outside the container image, so they remain available after the job ends.

6. Resume from a checkpoint

To resume training, submit the same Slurm script again:

sbatch example-04-checkpoint.sh

The submission script passes --resume-training to the Python script. If checkpoints already exist, the Python script loads the latest one and continues from the next epoch. The output file should include lines similar to:

Resuming from checkpoint: /users/abc123/tmp/CoSTARexamples/Example04/models_tf/ckpt-4 (start at epoch 5)
Train Epoch: 5 [    0]  Loss: 1.304010
Train Epoch: 5 [  320]  Loss: 1.466132

If no checkpoint exists, the script starts training from epoch 1.