Summary

The following document will detail on a high level how to use the py-tensorflow module on Compute2 to train a model. In many cases this assumes the use of the py-keras module as well.

All examples were derived from the official Tensorflow documentation here:

• https://www.tensorflow.org/tutorials

• https://www.tensorflow.org/tutorials/distribute/multi_worker_with_keras

• https://www.tensorflow.org/tutorials/distribute/keras

For more detailed information, examples, and How-Tos please consult the official documentation.

• https://www.tensorflow.org/api_docs/python/tf

Quick-Start

Single-Node

Single-GPU

Setup

Python Program

1. Create the following sample program as tf-1node-1gpu.py in your working directory.

   #!/usr/bin/env python3 ''' Import TensorFlow into your program: ''' import tensorflow as tf print("TensorFlow version:", tf.__version__) from tensorflow.keras.layers import Dense, Flatten, Conv2D from tensorflow.keras import Model ''' Load and prepare the MNIST dataset. ''' mnist = tf.keras.datasets.mnist (x_train, y_train), (x_test, y_test) = mnist.load_data() x_train, x_test = x_train / 255.0, x_test / 255.0 # Add a channels dimension x_train = x_train[..., tf.newaxis].astype("float32") x_test = x_test[..., tf.newaxis].astype("float32") ''' Use tf.data to batch and shuffle the dataset: ''' train_ds = tf.data.Dataset.from_tensor_slices( (x_train, y_train)).shuffle(10000).batch(32) test_ds = tf.data.Dataset.from_tensor_slices((x_test, y_test)).batch(32) ''' Build the tf.keras model using the Keras model subclassing API: ''' class MyModel(Model): def __init__(self): super().__init__() self.conv1 = Conv2D(32, 3, activation='relu') self.flatten = Flatten() self.d1 = Dense(128, activation='relu') self.d2 = Dense(10) def call(self, x): x = self.conv1(x) x = self.flatten(x) x = self.d1(x) return self.d2(x) # Create an instance of the model model = MyModel() ''' Choose an optimizer and loss function for training: ''' loss_object = tf.keras.losses.SparseCategoricalCrossentropy(from_logits=True) optimizer = tf.keras.optimizers.Adam() ''' Select metrics to measure the loss and the accuracy of the model. These metrics accumulate the values over epochs and then print the overall result. ''' train_loss = tf.keras.metrics.Mean(name='train_loss') train_accuracy = tf.keras.metrics.SparseCategoricalAccuracy(name='train_accuracy') test_loss = tf.keras.metrics.Mean(name='test_loss') test_accuracy = tf.keras.metrics.SparseCategoricalAccuracy(name='test_accuracy') ''' Use tf.GradientTape to train the model: ''' @tf.function def train_step(images, labels): with tf.GradientTape() as tape: # training=True is only needed if there are layers with different # behavior during training versus inference (e.g. Dropout). predictions = model(images, training=True) loss = loss_object(labels, predictions) gradients = tape.gradient(loss, model.trainable_variables) optimizer.apply_gradients(zip(gradients, model.trainable_variables)) train_loss(loss) train_accuracy(labels, predictions) ''' Test the model: ''' @tf.function def test_step(images, labels): # training=False is only needed if there are layers with different # behavior during training versus inference (e.g. Dropout). predictions = model(images, training=False) t_loss = loss_object(labels, predictions) test_loss(t_loss) test_accuracy(labels, predictions) EPOCHS = 5 for epoch in range(EPOCHS): # Reset the metrics at the start of the next epoch train_loss.reset_state() train_accuracy.reset_state() test_loss.reset_state() test_accuracy.reset_state() for images, labels in train_ds: train_step(images, labels) for test_images, test_labels in test_ds: test_step(test_images, test_labels) print( f'Epoch {epoch + 1}, ' f'Loss: {train_loss.result():0.2f}, ' f'Accuracy: {train_accuracy.result() * 100:0.2f}, ' f'Test Loss: {test_loss.result():0.2f}, ' f'Test Accuracy: {test_accuracy.result() * 100:0.2f}' )

Execution - Baremetal

srun

1. Launch an srun job directly on the host.

   srun \ -A compute2-account \ --partition=general-short \ --gpus=1 \ --pty bash

2. Load the Tensorflow and Keras modules

   module load ris module load py-tensorflow/2.17.1 py-keras/3.6.0

3. Execute the tf-1node-1gpu.py file using python3

   python3 ./tf-1node-1gpu.py
   [gunnar@c2-gpu-001 tensorflow]$ python3 ./test.py 2025-10-15 15:51:40.223031: I tensorflow/core/util/port.cc:153] oneDNN custom operations are on. You may see slightly different numerical results due to floating-point round-off errors from different computation orders. To turn them off, set the environment variable `TF_ENABLE_ONEDNN_OPTS=0`. 2025-10-15 15:51:40.235248: E external/local_xla/xla/stream_executor/cuda/cuda_fft.cc:485] Unable to register cuFFT factory: Attempting to register factory for plugin cuFFT when one has already been registered 2025-10-15 15:51:40.249130: E external/local_xla/xla/stream_executor/cuda/cuda_dnn.cc:8454] Unable to register cuDNN factory: Attempting to register factory for plugin cuDNN when one has already been registered 2025-10-15 15:51:40.253262: E external/local_xla/xla/stream_executor/cuda/cuda_blas.cc:1452] Unable to register cuBLAS factory: Attempting to register factory for plugin cuBLAS when one has already been registered 2025-10-15 15:51:40.262718: I tensorflow/core/platform/cpu_feature_guard.cc:210] This TensorFlow binary is optimized to use available CPU instructions in performance-critical operations. To enable the following instructions: AVX512_VNNI AVX512_BF16 AVX512_FP16 AVX_VNNI AMX_TILE AMX_INT8 AMX_BF16, in other operations, rebuild TensorFlow with the appropriate compiler flags. TensorFlow version: 2.17.1 2025-10-15 15:51:47.886970: I tensorflow/core/common_runtime/gpu/gpu_device.cc:2021] Created device /job:localhost/replica:0/task:0/device:GPU:0 with 8077 MB memory: -> device: 0, name: NVIDIA H100 80GB HBM3 MIG 1g.10gb, pci bus id: 0000:55:00.0, compute capability: 9.0 2025-10-15 15:51:50.148849: I external/local_xla/xla/stream_executor/cuda/cuda_dnn.cc:531] Loaded cuDNN version 8907 WARNING: All log messages before absl::InitializeLog() is called are written to STDERR W0000 00:00:1760561510.198783 1181457 gpu_timer.cc:114] Skipping the delay kernel, measurement accuracy will be reduced W0000 00:00:1760561510.223141 1181457 gpu_timer.cc:114] Skipping the delay kernel, measurement accuracy will be reduced W0000 00:00:1760561510.227272 1181457 gpu_timer.cc:114] Skipping the delay kernel, measurement accuracy will be reduced ... ... W0000 00:00:1760561514.110248 1181457 gpu_timer.cc:114] Skipping the delay kernel, measurement accuracy will be reduced W0000 00:00:1760561514.111084 1181457 gpu_timer.cc:114] Skipping the delay kernel, measurement accuracy will be reduced 2025-10-15 15:51:54.113154: I tensorflow/core/framework/local_rendezvous.cc:404] Local rendezvous is aborting with status: OUT_OF_RANGE: End of sequence Epoch 1, Loss: 0.14, Accuracy: 95.76, Test Loss: 0.07, Test Accuracy: 97.73 2025-10-15 15:51:57.396113: I tensorflow/core/framework/local_rendezvous.cc:404] Local rendezvous is aborting with status: OUT_OF_RANGE: End of sequence Epoch 2, Loss: 0.04, Accuracy: 98.66, Test Loss: 0.05, Test Accuracy: 98.34 Epoch 3, Loss: 0.02, Accuracy: 99.29, Test Loss: 0.06, Test Accuracy: 98.26 2025-10-15 15:52:03.981491: I tensorflow/core/framework/local_rendezvous.cc:404] Local rendezvous is aborting with status: OUT_OF_RANGE: End of sequence Epoch 4, Loss: 0.01, Accuracy: 99.59, Test Loss: 0.06, Test Accuracy: 98.53 Epoch 5, Loss: 0.01, Accuracy: 99.69, Test Loss: 0.07, Test Accuracy: 98.37

sbatch

1. Create an sbatch file named tf-1node-1gpu-host.sbatch for execution on a single host and save it to your working directory.

   #!/bin/bash #SBATCH -A compute2-account #SBATCH --partition=general-short #SBATCH --nodes=1 #SBATCH --gpus=1 # load modules module load ris module load py-tensorflow/2.17.1 py-keras/3.6.0 # Execute the test script srun python3 tf-1node-1gpu.py

2. Execute the sbatch job

   sbatch ./tf-1node-1gpu-host.sbatch
   [gunnar@c2-login-002 tensorflow]$ sbatch ./tf-1node-1gpu-host.sbatch Submitted batch job 163601

3. You can view the output of the job like so

   tail -f slurm-<job_id>.out
   [gunnar@c2-login-002 tensorflow]$ tail -f 163601.out W0000 00:00:1760714275.281526 2007089 gpu_timer.cc:114] Skipping the delay kernel, measurement accuracy will be reduced 2025-10-17 10:17:55.283653: I tensorflow/core/framework/local_rendezvous.cc:404] Local rendezvous is aborting with status: OUT_OF_RANGE: End of sequence 2025-10-17 10:17:58.593164: I tensorflow/core/framework/local_rendezvous.cc:404] Local rendezvous is aborting with status: OUT_OF_RANGE: End of sequence 2025-10-17 10:18:05.252018: I tensorflow/core/framework/local_rendezvous.cc:404] Local rendezvous is aborting with status: OUT_OF_RANGE: End of sequence TensorFlow version: 2.17.1 Epoch 1, Loss: 0.13, Accuracy: 95.93, Test Loss: 0.06, Test Accuracy: 98.15 Epoch 2, Loss: 0.04, Accuracy: 98.69, Test Loss: 0.05, Test Accuracy: 98.26 Epoch 3, Loss: 0.02, Accuracy: 99.26, Test Loss: 0.06, Test Accuracy: 98.32 Epoch 4, Loss: 0.01, Accuracy: 99.57, Test Loss: 0.06, Test Accuracy: 98.32 Epoch 5, Loss: 0.01, Accuracy: 99.71, Test Loss: 0.06, Test Accuracy: 98.38

Execution - Container

srun

1. Launch an srun job using the C2-THPC container.

   default_mounts='/etc/profile.d,/etc/sysconfig/modules,/opt/thpc,/storage2/fs1,/cm,/lib64/libmunge.so.2,/run/munge,/storage2/fs1,/scratch2/fs1,/storage1/fs1,/rdcw/fs1,/rdcw/fs2' NVIDIA_VISIBLE_DEVICES=all \ NVIDIA_DRIVER_CAPABILITIES=compute,graphics,utility \ srun \ -A compute2-account --partition=general-short \ --gpus=1 \ --container-image=ghcr.io#washu-it-ris/ris-thpc:rocky9.2 \ --container-mounts="$(pwd),${default_mounts}" \ --container-workdir=$(pwd) \ --pty bash

2. Load the Tensorflow and Keras modules

   module load ris module load py-tensorflow/2.17.1 py-keras/3.6.0

3. Execute the tf-1node-1gpu.py file using python3

   python3 ./tf-1node-1gpu.py
   [gunnar@c2-gpu-001 tensorflow]$ python3 ./test.py 2025-10-15 15:51:40.223031: I tensorflow/core/util/port.cc:153] oneDNN custom operations are on. You may see slightly different numerical results due to floating-point round-off errors from different computation orders. To turn them off, set the environment variable `TF_ENABLE_ONEDNN_OPTS=0`. 2025-10-15 15:51:40.235248: E external/local_xla/xla/stream_executor/cuda/cuda_fft.cc:485] Unable to register cuFFT factory: Attempting to register factory for plugin cuFFT when one has already been registered 2025-10-15 15:51:40.249130: E external/local_xla/xla/stream_executor/cuda/cuda_dnn.cc:8454] Unable to register cuDNN factory: Attempting to register factory for plugin cuDNN when one has already been registered 2025-10-15 15:51:40.253262: E external/local_xla/xla/stream_executor/cuda/cuda_blas.cc:1452] Unable to register cuBLAS factory: Attempting to register factory for plugin cuBLAS when one has already been registered 2025-10-15 15:51:40.262718: I tensorflow/core/platform/cpu_feature_guard.cc:210] This TensorFlow binary is optimized to use available CPU instructions in performance-critical operations. To enable the following instructions: AVX512_VNNI AVX512_BF16 AVX512_FP16 AVX_VNNI AMX_TILE AMX_INT8 AMX_BF16, in other operations, rebuild TensorFlow with the appropriate compiler flags. TensorFlow version: 2.17.1 2025-10-15 15:51:47.886970: I tensorflow/core/common_runtime/gpu/gpu_device.cc:2021] Created device /job:localhost/replica:0/task:0/device:GPU:0 with 8077 MB memory: -> device: 0, name: NVIDIA H100 80GB HBM3 MIG 1g.10gb, pci bus id: 0000:55:00.0, compute capability: 9.0 2025-10-15 15:51:50.148849: I external/local_xla/xla/stream_executor/cuda/cuda_dnn.cc:531] Loaded cuDNN version 8907 WARNING: All log messages before absl::InitializeLog() is called are written to STDERR W0000 00:00:1760561510.198783 1181457 gpu_timer.cc:114] Skipping the delay kernel, measurement accuracy will be reduced W0000 00:00:1760561510.223141 1181457 gpu_timer.cc:114] Skipping the delay kernel, measurement accuracy will be reduced W0000 00:00:1760561510.227272 1181457 gpu_timer.cc:114] Skipping the delay kernel, measurement accuracy will be reduced ... ... W0000 00:00:1760561514.110248 1181457 gpu_timer.cc:114] Skipping the delay kernel, measurement accuracy will be reduced W0000 00:00:1760561514.111084 1181457 gpu_timer.cc:114] Skipping the delay kernel, measurement accuracy will be reduced 2025-10-15 15:51:54.113154: I tensorflow/core/framework/local_rendezvous.cc:404] Local rendezvous is aborting with status: OUT_OF_RANGE: End of sequence Epoch 1, Loss: 0.14, Accuracy: 95.76, Test Loss: 0.07, Test Accuracy: 97.73 2025-10-15 15:51:57.396113: I tensorflow/core/framework/local_rendezvous.cc:404] Local rendezvous is aborting with status: OUT_OF_RANGE: End of sequence Epoch 2, Loss: 0.04, Accuracy: 98.66, Test Loss: 0.05, Test Accuracy: 98.34 Epoch 3, Loss: 0.02, Accuracy: 99.29, Test Loss: 0.06, Test Accuracy: 98.26 2025-10-15 15:52:03.981491: I tensorflow/core/framework/local_rendezvous.cc:404] Local rendezvous is aborting with status: OUT_OF_RANGE: End of sequence Epoch 4, Loss: 0.01, Accuracy: 99.59, Test Loss: 0.06, Test Accuracy: 98.53 Epoch 5, Loss: 0.01, Accuracy: 99.69, Test Loss: 0.07, Test Accuracy: 98.37

sbatch

1. Create an sbatch file named tf-1node-1gpu-container.sbatch to run our python program within the C2-THPC container and save it to your working directory.

   #!/bin/bash #SBATCH -A compute2-account #SBATCH --partition=general-short #SBATCH --nodes=1 #SBATCH --gpus=1 #SBATCH --container-image='ghcr.io#washu-it-ris/ris-thpc:rocky9.2' #SBATCH --container-mounts='/etc/profile.d,/etc/sysconfig/modules,/opt/thpc,/storage2/fs1,/cm,/lib64/libmunge.so.2,/run/munge,/storage2/fs1,/scratch2/fs1,/storage1/fs1,/rdcw/fs1,/rdcw/fs2' # Move to working-directory cd $SLURM_SUBMIT_DIR # Export NVIDIA variables for GPU access within the container export NVIDIA_VISIBLE_DEVICES=all export NVIDIA_DRIVER_CAPABILITIES=compute,utility,graphics # load modules module load ris module load py-tensorflow/2.17.1 py-keras/3.6.0 # Execute the test script python3 tf-1node-1gpu.py

2. Execute the sbatch job

   sbatch ./tf-1node-1gpu-container.batch
   [gunnar@c2-login-002 tensorflow]$ sbatch ./tf-1node-1gpu-container.sbatch Submitted batch job 163608

3. You can view the output of the job like so

   tail -f slurm-<job_id>.out
   [gunnar@c2-login-002 tensorflow]$ tail -f 163608.out 2025-10-17 10:55:34.043431: I tensorflow/core/framework/local_rendezvous.cc:404] Local rendezvous is aborting with status: OUT_OF_RANGE: End of sequence 2025-10-17 10:55:34.904181: I tensorflow/core/framework/local_rendezvous.cc:404] Local rendezvous is aborting with status: OUT_OF_RANGE: End of sequence 2025-10-17 10:56:09.006432: I tensorflow/core/framework/local_rendezvous.cc:404] Local rendezvous is aborting with status: OUT_OF_RANGE: End of sequence 2025-10-17 10:57:16.198291: I tensorflow/core/framework/local_rendezvous.cc:404] Local rendezvous is aborting with status: OUT_OF_RANGE: End of sequence TensorFlow version: 2.17.1 Epoch 1, Loss: 0.13, Accuracy: 95.98, Test Loss: 0.07, Test Accuracy: 97.76 Epoch 2, Loss: 0.04, Accuracy: 98.69, Test Loss: 0.07, Test Accuracy: 97.78 Epoch 3, Loss: 0.02, Accuracy: 99.25, Test Loss: 0.05, Test Accuracy: 98.38 Epoch 4, Loss: 0.01, Accuracy: 99.54, Test Loss: 0.06, Test Accuracy: 98.46 Epoch 5, Loss: 0.01, Accuracy: 99.64, Test Loss: 0.06, Test Accuracy: 98.32

ood

1. Launch a C2-OOD RIS Desktop session with a single GPU

   Open

   

2. Connect to the running C2-THPC session now running in C2-OOD

3. Load the Tensorflow and Keras modules

   module load ris module load py-tensorflow/2.17.1 py-keras/3.6.0
   Open

   

4. cd to your working directory with the tf-1node-1gpu.py file

5. Execute the tf-1node-1gpu.py file using python3

   python3 ./tf-1node-1gpu.py
   Open

   

Multi-GPU

Setup

Python Program

1. Save the following sample program as tf-1node-ngpu.py in your current working directory

   #!/usr/bin/env python3 ''' Setup ''' import tensorflow_datasets as tfds import tensorflow as tf import os # Load the TensorBoard notebook extension. #%load_ext tensorboard print(tf.__version__) ''' Download the dataset Load the MNIST dataset from TensorFlow Datasets. This returns a dataset in the tf.data format. Setting the with_info argument to True includes the metadata for the entire dataset, which is being saved here to info. Among other things, this metadata object includes the number of train and test examples. ''' datasets, info = tfds.load(name='mnist', with_info=True, as_supervised=True) mnist_train, mnist_test = datasets['train'], datasets['test'] ''' Define the distribution strategy Create a MirroredStrategy object. This will handle distribution and provide a context manager (MirroredStrategy.scope) to build your model inside. ''' strategy = tf.distribute.MirroredStrategy() print('Number of devices: {}'.format(strategy.num_replicas_in_sync)) ''' Set up the input pipeline When training a model with multiple GPUs, you can use the extra computing power effectively by increasing the batch size. In general, use the largest batch size that fits the GPU memory and tune the learning rate accordingly. ''' # You can also do info.splits.total_num_examples to get the total # number of examples in the dataset. num_train_examples = info.splits['train'].num_examples num_test_examples = info.splits['test'].num_examples BUFFER_SIZE = 10000 BATCH_SIZE_PER_REPLICA = 64 BATCH_SIZE = BATCH_SIZE_PER_REPLICA * strategy.num_replicas_in_sync ''' Define a function that normalizes the image pixel values from the [0, 255] range to the [0, 1] range (feature scaling): ''' def scale(image, label): image = tf.cast(image, tf.float32) image /= 255 return image, label ''' Apply this scale function to the training and test data, and then use the tf.data.Dataset APIs to shuffle the training data (Dataset.shuffle), and batch it (Dataset.batch). Notice that you are also keeping an in-memory cache of the training data to improve performance (Dataset.cache). ''' train_dataset = mnist_train.map(scale).cache().shuffle(BUFFER_SIZE).batch(BATCH_SIZE) eval_dataset = mnist_test.map(scale).batch(BATCH_SIZE) ''' Create the model and instantiate the optimizer Within the context of Strategy.scope, create and compile the model using the Keras API: ''' with strategy.scope(): model = tf.keras.Sequential([ tf.keras.layers.Conv2D(32, 3, activation='relu', input_shape=(28, 28, 1)), tf.keras.layers.MaxPooling2D(), tf.keras.layers.Flatten(), tf.keras.layers.Dense(64, activation='relu'), tf.keras.layers.Dense(10) ]) model.compile(loss=tf.keras.losses.SparseCategoricalCrossentropy(from_logits=True), optimizer=tf.keras.optimizers.Adam(learning_rate=0.001), metrics=['accuracy']) ''' Define the callbacks Define the following Keras Callbacks: tf.keras.callbacks.TensorBoard: writes a log for TensorBoard, which allows you to visualize the graphs. tf.keras.callbacks.ModelCheckpoint: saves the model at a certain frequency, such as after every epoch. tf.keras.callbacks.BackupAndRestore: provides the fault tolerance functionality by backing up the model and current epoch number. Learn more in the Fault tolerance section of the Multi-worker training with Keras tutorial. tf.keras.callbacks.LearningRateScheduler: schedules the learning rate to change after, for example, every epoch/batch. For illustrative purposes, add a custom callback called PrintLR to display the learning rate in the notebook. ''' # Define the checkpoint directory to store the checkpoints. checkpoint_dir = './training_checkpoints' # Define the name of the checkpoint files. #checkpoint_prefix = os.path.join(checkpoint_dir, "ckpt_{epoch}") checkpoint_prefix = os.path.join(checkpoint_dir, "ckpt_{epoch}.weights.h5") # Define a function for decaying the learning rate. # You can define any decay function you need. def decay(epoch): if epoch < 3: return 1e-3 elif epoch >= 3 and epoch < 7: return 1e-4 else: return 1e-5 # Define a callback for printing the learning rate at the end of each epoch. class PrintLR(tf.keras.callbacks.Callback): def on_epoch_end(self, epoch, logs=None): #print('\nLearning rate for epoch {} is {}'.format( epoch + 1, model.optimizer.lr.numpy())) print('\nLearning rate for epoch {} is {}'.format( epoch + 1, model.optimizer.learning_rate.numpy())) # Put all the callbacks together. callbacks = [ tf.keras.callbacks.TensorBoard(log_dir='./logs'), tf.keras.callbacks.ModelCheckpoint(filepath=checkpoint_prefix, save_weights_only=True), tf.keras.callbacks.LearningRateScheduler(decay), PrintLR() ] ''' Train and evaluate Now, train the model in the usual way by calling Keras Model.fit on the model and passing in the dataset created at the beginning of the tutorial. This step is the same whether you are distributing the training or not. ''' EPOCHS = 12 model.fit(train_dataset, epochs=EPOCHS, callbacks=callbacks) ''' Check for saved checkpoints: $ # Check the checkpoint directory. $ ls {checkpoint_dir} ''' ''' To check how well the model performs, load the latest checkpoint and call Model.evaluate on the test data: ''' #model.load_weights(tf.train.latest_checkpoint(checkpoint_dir)) files = os.listdir(checkpoint_dir) checkpoint_files = [os.path.join(checkpoint_dir, f) for f in files if f.endswith('.weights.h5')] if checkpoint_files: latest_checkpoint = max(checkpoint_files, key=os.path.getctime) model.load_weights(latest_checkpoint) else: print("No checkpoint files found.") eval_loss, eval_acc = model.evaluate(eval_dataset) print('Eval loss: {}, Eval accuracy: {}'.format(eval_loss, eval_acc)) ''' To visualize the output, launch TensorBoard and view the logs: %tensorboard --logdir=logs $ ls -sh ./logs ''' ''' Save the model Save the model to a .keras zip archive using Model.save. After your model is saved, you can load it with or without the Strategy.scope. ''' path = 'my_model.keras' model.save(path) print('Saved model to: {}'.format(path)) ''' Now, load the model without Strategy.scope: ''' unreplicated_model = tf.keras.models.load_model(path) unreplicated_model.compile( loss=tf.keras.losses.SparseCategoricalCrossentropy(from_logits=True), optimizer=tf.keras.optimizers.Adam(), metrics=['accuracy']) eval_loss, eval_acc = unreplicated_model.evaluate(eval_dataset) print('Eval loss: {}, Eval Accuracy: {}'.format(eval_loss, eval_acc)) ''' Load the model with Strategy.scope: ''' with strategy.scope(): replicated_model = tf.keras.models.load_model(path) replicated_model.compile(loss=tf.keras.losses.SparseCategoricalCrossentropy(from_logits=True), optimizer=tf.keras.optimizers.Adam(), metrics=['accuracy']) eval_loss, eval_acc = replicated_model.evaluate(eval_dataset) print ('Eval loss: {}, Eval Accuracy: {}'.format(eval_loss, eval_acc))

Conda Environment

1. Create the following file as environment.yaml in your current working directory

   name: tensorflow-extras channels: - defaults dependencies: - python=3.11 - tensorflow-datasets

2. Install the conda-env named tensorflow-extras

   1. Load anaconda3 module

      module load ris module load anaconda3/2023.09-0

   2. Activate the base environment

      source activate

   3. Install the tensorflow-extras environment

      conda-env create -f environment.yaml
      # or update if it already exists conda-env update -f environment.yaml

   4. Deactivate the base environment

      conda deactivate

3. Identify the path to the tensorflow-extras environment’s site-package dir.

   1. Load anaconda3 module

      module load ris module load anaconda3/2023.09-0

   2. Find the path to the environment’s site-packages. This will be appended to PYTHONPATH later.

      conda_env_path=$(conda-env list | grep tensorflow-extras | tail -n1 | awk '{print $2}') conda_env_pyver=$(grep python environment.yaml | awk -F'=' '{print $2}') export TENSORFLOW_EXTRAS_SITEPACKAGES="$conda_env_path/lib/python$conda_env_pyver/site-packages" echo $TENSORFLOW_EXTRAS_SITEPACKAGES

Execution - Baremetal

srun

1. Launch an srun job directly on the host. Up to 6 GPUs can be used depending on available hosts.

   srun \ -A compute2-account \ --partition=general-gpu \ --gpus=2 \ --pty bash

2. Load the Tensorflow and Keras modules

   module load ris module load py-tensorflow/2.17.1 py-keras/3.6.0

3. Append the path to site-packages in the tensorflow-extras conda environment to the PYTHONPATH environment variable.

   export PYTHONPATH=$PYTHONPATH:$TENSORFLOW_EXTRAS_SITEPACKAGES

4. Execute the tf-1node-ngpu.py file using python3

   python3 ./tf-1node-ngpu.py
   [gunnar@c2-gpu-001 tensorflow]$ python3 ./test.py 2025-10-15 15:51:40.223031: I tensorflow/core/util/port.cc:153] oneDNN custom operations are on. You may see slightly different numerical results due to floating-point round-off errors from different computation orders. To turn them off, set the environment variable `TF_ENABLE_ONEDNN_OPTS=0`. 2025-10-15 15:51:40.235248: E external/local_xla/xla/stream_executor/cuda/cuda_fft.cc:485] Unable to register cuFFT factory: Attempting to register factory for plugin cuFFT when one has already been registered 2025-10-15 15:51:40.249130: E external/local_xla/xla/stream_executor/cuda/cuda_dnn.cc:8454] Unable to register cuDNN factory: Attempting to register factory for plugin cuDNN when one has already been registered 2025-10-15 15:51:40.253262: E external/local_xla/xla/stream_executor/cuda/cuda_blas.cc:1452] Unable to register cuBLAS factory: Attempting to register factory for plugin cuBLAS when one has already been registered 2025-10-15 15:51:40.262718: I tensorflow/core/platform/cpu_feature_guard.cc:210] This TensorFlow binary is optimized to use available CPU instructions in performance-critical operations. To enable the following instructions: AVX512_VNNI AVX512_BF16 AVX512_FP16 AVX_VNNI AMX_TILE AMX_INT8 AMX_BF16, in other operations, rebuild TensorFlow with the appropriate compiler flags. TensorFlow version: 2.17.1 2025-10-15 15:51:47.886970: I tensorflow/core/common_runtime/gpu/gpu_device.cc:2021] Created device /job:localhost/replica:0/task:0/device:GPU:0 with 8077 MB memory: -> device: 0, name: NVIDIA H100 80GB HBM3 MIG 1g.10gb, pci bus id: 0000:55:00.0, compute capability: 9.0 2025-10-15 15:51:50.148849: I external/local_xla/xla/stream_executor/cuda/cuda_dnn.cc:531] Loaded cuDNN version 8907 WARNING: All log messages before absl::InitializeLog() is called are written to STDERR W0000 00:00:1760561510.198783 1181457 gpu_timer.cc:114] Skipping the delay kernel, measurement accuracy will be reduced W0000 00:00:1760561510.223141 1181457 gpu_timer.cc:114] Skipping the delay kernel, measurement accuracy will be reduced W0000 00:00:1760561510.227272 1181457 gpu_timer.cc:114] Skipping the delay kernel, measurement accuracy will be reduced ... ... W0000 00:00:1760561514.110248 1181457 gpu_timer.cc:114] Skipping the delay kernel, measurement accuracy will be reduced W0000 00:00:1760561514.111084 1181457 gpu_timer.cc:114] Skipping the delay kernel, measurement accuracy will be reduced 2025-10-15 15:51:54.113154: I tensorflow/core/framework/local_rendezvous.cc:404] Local rendezvous is aborting with status: OUT_OF_RANGE: End of sequence Epoch 1, Loss: 0.14, Accuracy: 95.76, Test Loss: 0.07, Test Accuracy: 97.73 2025-10-15 15:51:57.396113: I tensorflow/core/framework/local_rendezvous.cc:404] Local rendezvous is aborting with status: OUT_OF_RANGE: End of sequence Epoch 2, Loss: 0.04, Accuracy: 98.66, Test Loss: 0.05, Test Accuracy: 98.34 Epoch 3, Loss: 0.02, Accuracy: 99.29, Test Loss: 0.06, Test Accuracy: 98.26 2025-10-15 15:52:03.981491: I tensorflow/core/framework/local_rendezvous.cc:404] Local rendezvous is aborting with status: OUT_OF_RANGE: End of sequence Epoch 4, Loss: 0.01, Accuracy: 99.59, Test Loss: 0.06, Test Accuracy: 98.53 Epoch 5, Loss: 0.01, Accuracy: 99.69, Test Loss: 0.07, Test Accuracy: 98.37

5. You should now also have a logs dir, a training_checkpoints dir, and a my_model.keras file

   [gunnar@c2-gpu-001 tensorflow]$ ls -al total 4298 drwx--S---+ 2 gunnar compute2-ris 4096 Oct 21 09:57 . drwx--S---+ 2 gunnar compute2-ris 4096 Oct 15 15:17 .. -rw-------+ 1 gunnar compute2-ris 102 Oct 17 12:18 environment.yaml drwx--S---+ 2 gunnar compute2-ris 4096 Oct 21 09:56 logs -rw-------+ 1 gunnar compute2-ris 4195975 Oct 21 09:57 my_model.keras -rw-------+ 1 gunnar compute2-ris 7090 Oct 17 15:53 tf-1node-ngpu.py drwx--S---+ 2 gunnar compute2-ris 4096 Oct 21 09:57 training_checkpoints

sbatch

1. Create an sbatch file named tf-1node-ngpu-host.sbatch for execution on a single host and save it to your working directory.

   #!/bin/bash #SBATCH -A compute2-account #SBATCH --partition=general-gpu #SBATCH --nodes=1 #SBATCH --gpus=2 # load modules module load ris module load py-tensorflow/2.17.1 py-keras/3.6.0 # add `tensorflow-extras` conda-env to PYTHONPATH export PYTHONPATH="$PYTHONPATH:$TENSORFLOW_EXTRAS_SITEPACKAGES" # Execute the test script python3 tf-1node-ngpu.py

2. Execute the sbatch job

   sbatch ./tf-1node-ngpu-host.sbatch
   [gunnar@c2-login-002 tensorflow]$ sbatch tf-1node-ngpu-host.sbatch Submitted batch job 164926

3. You can view the output of the job like so

   tail -f slurm-<job_id>.out
   [gunnar@c2-login-002 tensorflow]$ tail -f slurm-164926.out 2025-10-21 09:57:27.524188: W external/local_xla/xla/service/gpu/gemm_fusion_autotuner.cc:806] Compiling 22 configs for gemm_fusion_dot.50 on a single thread. I0000 00:00:1761058648.967071 511984 device_compiler.h:188] Compiled cluster using XLA! This line is logged at most once for the lifetime of the process. Eval loss: 0.03913751244544983, Eval accuracy: 0.9848242998123169 Saved model to: my_model.keras 63/79 ━━━━━━━━━━━━━━━━━━━━ 0s 2ms/step - accuracy: 0.9863 - loss: 0.04002025-10-21 09:57:29.189795: W external/local_xla/xla/service/gpu/gemm_fusion_autotuner.cc:806] Compiling 2 configs for gemm_fusion_dot.50 on a single thread. 79/79 ━━━━━━━━━━━━━━━━━━━━ 2s 7ms/step - accuracy: 0.9862 - loss: 0.0404 2025-10-21 09:57:29.616629: W tensorflow/core/grappler/optimizers/data/auto_shard.cc:553] The `assert_cardinality` transformation is currently not handled by the auto-shard rewrite and will be removed. Eval loss: 0.041532840579748154, Eval Accuracy: 0.9858999848365784 79/79 ━━━━━━━━━━━━━━━━━━━━ 1s 5ms/step - accuracy: 0.9861 - loss: nan Eval loss: 0.03913751244544983, Eval Accuracy: 0.9848242998123169

4. You should now also have a logs dir, a training_checkpoints dir, and a my_model.keras file

   [gunnar@c2-login-002 tensorflow]$ ls -al total 4298 drwx--S---+ 2 gunnar compute2-ris 4096 Oct 21 09:57 . drwx--S---+ 2 gunnar compute2-ris 4096 Oct 15 15:17 .. -rw-------+ 1 gunnar compute2-ris 102 Oct 17 12:18 environment.yaml drwx--S---+ 2 gunnar compute2-ris 4096 Oct 21 09:56 logs -rw-------+ 1 gunnar compute2-ris 4195975 Oct 21 09:57 my_model.keras -rw-------+ 1 gunnar compute2-ris 191079 Oct 21 09:57 slurm-164926.out -rw-------+ 1 gunnar compute2-ris 328 Oct 21 09:50 tf-1node-ngpu-host.sbatch -rw-------+ 1 gunnar compute2-ris 7090 Oct 17 15:53 tf-1node-ngpu.py drwx--S---+ 2 gunnar compute2-ris 4096 Oct 21 09:57 training_checkpoints

Execution - Container

srun

1. Launch an srun job using the C2-THPC container.

   default_mounts='/etc/profile.d,/etc/sysconfig/modules,/opt/thpc,/storage2/fs1,/cm,/lib64/libmunge.so.2,/run/munge,/storage2/fs1,/scratch2/fs1,/storage1/fs1,/rdcw/fs1,/rdcw/fs2' NVIDIA_VISIBLE_DEVICES=all \ NVIDIA_DRIVER_CAPABILITIES=compute,graphics,utility \ srun \ -A compute2-account \ --partition=general-gpu \ --gpus=2 \ --container-image=ghcr.io#washu-it-ris/ris-thpc:rocky9.2 \ --container-mounts="$(pwd),${default_mounts}" \ --container-workdir=$(pwd) \ --pty bash

2. Load the Tensorflow and Keras modules

   module load ris module load py-tensorflow/2.17.1 py-keras/3.6.0

3. Append the path to site-packages in the tensorflow-extras conda environment to the PYTHONPATH environment variable.

   export PYTHONPATH=$PYTHONPATH:$TENSORFLOW_EXTRAS_SITEPACKAGES

4. Execute the tf-1node-ngpu.py file using python3

   python3 ./tf-1node-ngpu.py
   [gunnar@c2-gpu-016 tensorflow]$ python3 ./tf-1node-ngpu.py 2025-10-21 10:14:58.251504: I tensorflow/core/util/port.cc:153] oneDNN custom operations are on. You may see slightly different numerical results due to floating-point round-off errors from different computation orders. To turn them of f, set the environment variable `TF_ENABLE_ONEDNN_OPTS=0`. 2025-10-21 10:14:58.263477: E external/local_xla/xla/stream_executor/cuda/cuda_fft.cc:485] Unable to register cuFFT factory: Attempting to register factory for plugin cuFFT when one has already been registered 2025-10-21 10:14:58.277184: E external/local_xla/xla/stream_executor/cuda/cuda_dnn.cc:8454] Unable to register cuDNN factory: Attempting to register factory for plugin cuDNN when one has already been registered 2025-10-21 10:14:58.281281: E external/local_xla/xla/stream_executor/cuda/cuda_blas.cc:1452] Unable to register cuBLAS factory: Attempting to register factory for plugin cuBLAS when one has already been registered 2025-10-21 10:14:58.290740: I tensorflow/core/platform/cpu_feature_guard.cc:210] This TensorFlow binary is optimized to use available CPU instructions in performance-critical operations. To enable the following instructions: AVX512_VNNI AVX512_BF16 AVX512_FP16 AVX_VNNI AMX_TILE AMX_INT8 AMX_BF16, in other operations, rebuild TensorFlow with the appropriate compiler flags. 2.17.1 2025-10-21 10:15:08.316460: I tensorflow/core/common_runtime/gpu/gpu_device.cc:2021] Created device /job:localhost/replica:0/task:0/device:GPU:0 with 78668 MB memory: -> device: 0, name: NVIDIA H100 80GB HBM3, pci bus id: 0000:55:00.0 , compute capability: 9.0 2025-10-21 10:15:08.318305: I tensorflow/core/common_runtime/gpu/gpu_device.cc:2021] Created device /job:localhost/replica:0/task:0/device:GPU:1 with 78668 MB memory: -> device: 1, name: NVIDIA H100 80GB HBM3, pci bus id: 0000:68:00.0 , compute capability: 9.0 Number of devices: 2 /opt/thpc/linux/x86_64/ris/software/gcc-13.3.0/py-keras-3.6.0-eu4zvyhrsx42im5qta44sbonk3evivif/lib/python3.11/site-packages/keras/src/layers/convolutional/base_conv.py:107: UserWarning: Do not pass an `input_shape`/`input_dim` argument to a layer. When using Sequential models, prefer using an `Input(shape)` object as the first layer in the model instead. super().__init__(activity_regularizer=activity_regularizer, **kwargs) 2025-10-21 10:15:10.297938: W tensorflow/core/grappler/optimizers/data/auto_shard.cc:553] The `assert_cardinality` transformation is currently not handled by the auto-shard rewrite and will be removed. 2025-10-21 10:15:11.058042: W tensorflow/core/kernels/data/cache_dataset_ops.cc:913] The calling iterator did not fully read the dataset being cached. In order to avoid unexpected truncation of the dataset, the partially cached content s of the dataset will be discarded. This can happen if you have an input pipeline similar to `dataset.cache().take(k).repeat()`. You should use `dataset.take(k).cache().repeat()` instead. 2025-10-21 10:15:11.058904: W tensorflow/core/kernels/data/cache_dataset_ops.cc:913] The calling iterator did not fully read the dataset being cached. In order to avoid unexpected truncation of the dataset, the partially cached content s of the dataset will be discarded. This can happen if you have an input pipeline similar to `dataset.cache().take(k).repeat()`. You should use `dataset.take(k).cache().repeat()` instead. Epoch 1/12 2025-10-21 10:15:14.480825: I external/local_xla/xla/stream_executor/cuda/cuda_dnn.cc:531] Loaded cuDNN version 8907 WARNING: All log messages before absl::InitializeLog() is called are written to STDERR W0000 00:00:1761059714.532818 517386 gpu_timer.cc:114] Skipping the delay kernel, measurement accuracy will be reduced W0000 00:00:1761059714.555719 517386 gpu_timer.cc:114] Skipping the delay kernel, measurement accuracy will be reduced ... ... W0000 00:00:1761059776.421183 517386 gpu_timer.cc:114] Skipping the delay kernel, measurement accuracy will be reduced W0000 00:00:1761059776.421376 517386 gpu_timer.cc:114] Skipping the delay kernel, measurement accuracy will be reduced 79/79 ━━━━━━━━━━━━━━━━━━━━ 1s 7ms/step - accuracy: 0.9845 - loss: nan Eval loss: 0.044944487512111664, Eval accuracy: 0.9832268357276917 Saved model to: my_model.keras WARNING: All log messages before absl::InitializeLog() is called are written to STDERR I0000 00:00:1761059776.720392 517386 service.cc:146] XLA service 0x7fd2730a63c0 initialized for platform CUDA (this does not guarantee that XLA will be used). Devices: I0000 00:00:1761059776.720425 517386 service.cc:154] StreamExecutor device (0): NVIDIA H100 80GB HBM3, Compute Capability 9.0 I0000 00:00:1761059776.720429 517386 service.cc:154] StreamExecutor device (1): NVIDIA H100 80GB HBM3, Compute Capability 9.0 2025-10-21 10:16:16.727701: I tensorflow/compiler/mlir/tensorflow/utils/dump_mlir_util.cc:268] disabling MLIR crash reproducer, set env var `MLIR_CRASH_REPRODUCER_DIRECTORY` to enable. 2025-10-21 10:16:16.749352: W external/local_xla/xla/service/gpu/gemm_fusion_autotuner.cc:806] Compiling 22 configs for gemm_fusion_dot.50 on a single thread. I0000 00:00:1761059778.108872 517386 device_compiler.h:188] Compiled cluster using XLA! This line is logged at most once for the lifetime of the process. 65/79 ━━━━━━━━━━━━━━━━━━━━ 0s 2ms/step - accuracy: 0.9846 - loss: 0.04472025-10-21 10:16:18.326256: W external/local_xla/xla/service/gpu/gemm_fusion_autotuner.cc:806] Compiling 2 configs for gemm_fusion_dot.50 on a single thread. 79/79 ━━━━━━━━━━━━━━━━━━━━ 2s 7ms/step - accuracy: 0.9845 - loss: 0.0450 Eval loss: 0.04601269215345383, Eval Accuracy: 0.984000027179718 2025-10-21 10:16:18.754512: W tensorflow/core/grappler/optimizers/data/auto_shard.cc:553] The `assert_cardinality` transformation is currently not handled by the auto-shard rewrite and will be removed. 79/79 ━━━━━━━━━━━━━━━━━━━━ 1s 4ms/step - accuracy: 0.9846 - loss: nan Eval loss: 0.044944487512111664, Eval Accuracy: 0.9832268357276917

5. You should now also have a logs dir, a training_checkpoints dir, and a my_model.keras file

   [gunnar@c2-gpu-016 tensorflow]$ ls -al total 4111 drwx--S---+ 2 gunnar nogroup 4096 Oct 21 10:16 . drwx--S---+ 2 gunnar nogroup 4096 Oct 15 15:17 .. -rw-------+ 1 gunnar nogroup 102 Oct 17 12:18 environment.yaml drwx--S---+ 2 gunnar nogroup 4096 Oct 21 10:15 logs -rw-------+ 1 gunnar nogroup 4195975 Oct 21 10:16 my_model.keras -rw-------+ 1 gunnar nogroup 7090 Oct 17 15:53 tf-1node-ngpu.py drwx--S---+ 2 gunnar nogroup 4096 Oct 21 10:16 training_checkpoints

sbatch

1. Identify your current working directory using the pwd command

   pwd
   [gunnar@c2-login-002 tensorflow]$ pwd /scratch2/fs1/ris/gunnar/tmp/tensorflow

2. Create an sbatch file named tf-1node-ngpu-container.sbatch to run our python program within the C2-THPC container and save it to your working directory.

   #!/bin/bash #SBATCH -A compute2-account #SBATCH --partition=general-gpu #SBATCH --nodes=1 #SBATCH --gpus=2 #SBATCH --container-image='ghcr.io#washu-it-ris/ris-thpc:rocky9.2' #SBATCH --container-mounts='/etc/profile.d,/etc/sysconfig/modules,/opt/thpc,/storage2/fs1,/cm,/lib64/libmunge.so.2,/run/munge,/storage2/fs1,/scratch2/fs1,/storage1/fs1,/rdcw/fs1,/rdcw/fs2' # Move to working-directory cd $SLURM_SUBMIT_DIR # Export NVIDIA variables for GPU access within the container export NVIDIA_VISIBLE_DEVICES=all export NVIDIA_DRIVER_CAPABILITIES=compute,utility,graphics # add `tensorflow-extras` conda-env to PYTHONPATH export PYTHONPATH="$PYTHONPATH:$TENSORFLOW_EXTRAS_SITEPACKAGES" # load modules module load ris module load py-tensorflow/2.17.1 py-keras/3.6.0 # Execute the test script python3 tf-1node-ngpu.py

3. Execute the sbatch job

   sbatch ./tf-1node-ngpu-container.sbatch
   [gunnar@c2-login-002 tensorflow]$ sbatch ./tf-1node-ngpu-container.sbatch Submitted batch job 165482

4. You can view the output of the job like so

   tail -f slurm-<job_id>.out
   [gunnar@c2-login-002 tensorflow]$ tail -f slurm-165482.out 2025-10-23 11:22:20.181503: W tensorflow/core/framework/dataset.cc:993] Input of GeneratorDatasetOp::Dataset will not be optimized because the dataset does not implement the AsGraphDefInternal() method needed to apply optimizations. 2025-10-23 11:22:20.208736: W tensorflow/core/kernels/data/cache_dataset_ops.cc:913] The calling iterator did not fully read the dataset being cached. In order to avoid unexpected truncation of the dataset, the partially cached contents of the dataset will be discarded. This can happen if you have an input pipeline similar to `dataset.cache().take(k).repeat()`. You should use `dataset.take(k).cache().repeat()` instead. 157/157 ━━━━━━━━━━━━━━━━━━━━ 1s 5ms/step - accuracy: 0.9887 - loss: 0.0375 Eval loss: 0.03846264258027077, Eval accuracy: 0.9872999787330627 Saved model to: my_model.keras 157/157 ━━━━━━━━━━━━━━━━━━━━ 1s 3ms/step - accuracy: 0.9887 - loss: 0.0375 2025-10-23 11:22:22.052337: W tensorflow/core/grappler/optimizers/data/auto_shard.cc:553] The `assert_cardinality` transformation is currently not handled by the auto-shard rewrite and will be removed. Eval loss: 0.03846264258027077, Eval Accuracy: 0.9872999787330627 157/157 ━━━━━━━━━━━━━━━━━━━━ 1s 3ms/step - accuracy: 0.9887 - loss: 0.0375 Eval loss: 0.03846264258027077, Eval Accuracy: 0.9872999787330627

Multi-Node

Multi-GPU

Execution - Baremetal

srun

sbatch

Execution - Container

srun

sbatch