Research IT

Tensorflow

TensorFlow is a free, open-source software library developed by the Google Brain Team used to build, train, and deploy machine learning and deep learning models. There are two ways to install Tensorflow in the cluster.  Either using micromamba or a Python virtual env.

Table of Contents

First we need to allocate a compute node to build our environment (change the arguments for your needs):

salloc -n 8 -N1 --mem=32G --gres=gpu:1 --time=8:00:00 

Installation Using micromamba

First load the micromamba module:

module load micromamba

Now we need to create a work directory:
micromamba create -prefix /path/to/work/directory/tf-mamba python=3.11 -c conda-forge -y

Changing /path/to/work/directory to your work directory.  Next hook the shell:

eval "$(micromamba shell hook --shell=bash)"

And activate your environment:

micromamba activate /path/to/work/directory/tf-mamba

There are two ways to install Tensorflow, with and without GPU support.

For CPU support:

pip install tensorflow

For GPU support:

pip install tensorflow[and-cuda]

Installation Using a Python env

First we load the correct python version (e.g.):
module load python/3.11.11-dnlm5qp

Then create your work venv"

python -m venv /path/to/work/directory/tf-venv

Then we activate the environment:

source /path/to/work/directory/tf-venv/bin/activate

Install Tensorflow

There are two ways to install Tensorflow, with and without GPU support.

For CPU support:

pip install tensorflow

For GPU support:

pip install tensorflow[and-cuda]

Testing the Environment

Below is a sample script you can use to test that Tensorflow is working.

import tensorflow as tf
import time

# 1. Check for GPU availability
print("TensorFlow Version:", tf.__version__)
gpus = tf.config.list_physical_devices('GPU')

if gpus:
    print(f"GPU is AVAILABLE. Found {len(gpus)} GPU(s):")
    for gpu in gpus:
        print(f"  - {gpu}")
    device_name = "/GPU:0"
else:
    print("GPU is NOT available. Falling back to CPU.")
    device_name = "/CPU:0"

# 2. Run a simple matrix multiplication test
print(f"Running calculation on {device_name}...")
with tf.device(device_name):
    size = 5000
    x = tf.random.normal([size, size])
    y = tf.random.normal([size, size])
    
    start_time = time.time()
    z = tf.matmul(x, y)
    
    # Force execution to finish to get accurate timing
    _ = z.numpy() 
    end_time = time.time()

print(f"Multiplication complete! Shape: {z.shape}")
print(f"Time taken: {end_time - start_time:.4f} seconds")

Save it as test_tf.py and run it:

python ./test_tf.py

If all is successful you should receive output that looks like this:

TensorFlow Version: 2.20.0
GPU is AVAILABLE. Found 1 GPU(s):
  - PhysicalDevice(name='/physical_device:GPU:0', device_type='GPU')
Running calculation on /GPU:0...
WARNING: All log messages before absl::InitializeLog() is called are written to STDERR
I0000 00:00:1787861904.596991 1276796 gpu_device.cc:2020] Created device /job:localhost/replica:0/task:0/device:GPU:0 with 79189 MB memory:  -> device: 0, name: NVIDIA A100-SXM4-80GB, pci bus id: 0000:81:00.0, compute capability: 8.0
Multiplication complete! Shape: (5000, 5000)
Time taken: 0.0828 seconds

Tutorials

Tensorflow has multiple tutorials available on their website.  They can be found here: