.. _tutorials-example-04-checkpoint:

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:

.. code-block:: bash

   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.

* :download:`example-04-checkpoint.sh <source_files/example-04-checkpoint.sh>`

.. dropdown:: example-04-checkpoint.sh (click to collapse / view)
   :color: primary
   :open:
   :icon: file-code

   .. literalinclude:: source_files/example-04-checkpoint.sh
      :language: bash
      :linenos:

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.

* :download:`example-04-checkpoint.py <source_files/example-04-checkpoint.py>`

.. dropdown:: example-04-checkpoint.py (click to collapse / view)
   :color: primary
   :open:
   :icon: file-code

   .. literalinclude:: source_files/example-04-checkpoint.py
      :language: python
      :linenos:

4. Submit your job
------------------

Submit the script to the Slurm scheduler:

.. code-block:: bash

   sbatch example-04-checkpoint.sh

You can check the status of your job by running:

.. code-block:: bash

   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:

.. code-block:: text

   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:

.. code-block:: bash

   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:

.. code-block:: text

   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.