Tensorflow

Tensorflow is a Python deep learning package from Google. The easiest way to use install it is to build a Python virtualenv with it in it.

First load GPU modules to allow access to accelerated GPGPU training.

module add cuda/11.8.0 cudnn/v8.6.0

Next you will want to create a virtualenv and source into it. Note that depending on the version of Tensorflow you need, you may also need to load a module for a more recent version of Python3.

$ python3 -m venv env
$ source env/bin/activate
(env) $

Then ensure you have a recent copy of pip in your virtualenv.

(env) $ pip install --upgrade pip
Collecting pip
  Downloading https://files.pythonhosted.org/packages/a4/6d/6463d49a933f547439d6b5b98b46af8742cc03ae83543e4d7688c2420f8b/pip-21.3.1-py3-none-any.whl (1.7MB)
    100% |████████████████████████████████| 1.7MB 1.2MB/s
Installing collected packages: pip
  Found existing installation: pip 9.0.3
    Uninstalling pip-9.0.3:
      Successfully uninstalled pip-9.0.3
Successfully installed pip-21.3.1

Then install the Tensorflow wheel through pip.

(env) $ pip install --upgrade tensorflow
Collecting tensorflow
  Downloading tensorflow-2.6.2-cp36-cp36m-manylinux2010_x86_64.whl (458.3 MB)
     |████████████████████████████████| 458.3 MB 37 kB/s
Collecting gast==0.4.0
  Downloading gast-0.4.0-py3-none-any.whl (9.8 kB)
...

Finally, start up a python shell (or install ipython through pip) and import Tensorflow.

(env)[username@hostname:/scratch0/username ] $ python
Python 3.6.8 (default, Oct  2 2023, 21:12:58)
[GCC 8.5.0 20210514 (Red Hat 8.5.0-18)] on linux
Type "help", "copyright", "credits" or "license" for more information.
>>> import tensorflow as tf
>>> tf.__version__
'2.6.2'

You can then try a more rigorous test by running the following example. Note that you may need to export XLA_FLAGS in your shell: export XLA_FLAGS=--xla_gpu_cuda_data_dir=/opt/common/cuda/cuda-x.x.x

import tensorflow as tf
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

model = tf.keras.models.Sequential([
  tf.keras.layers.Flatten(),
  tf.keras.layers.Dense(512, activation=tf.nn.relu),
  tf.keras.layers.Dropout(0.2),
  tf.keras.layers.Dense(10, activation=tf.nn.softmax)
])
model.compile(optimizer='adam',
              loss='sparse_categorical_crossentropy',
              metrics=['accuracy'])

model.fit(x_train, y_train, epochs=5)
model.evaluate(x_test, y_test)

To use this install after you close the shell you did this install in, you will need to both add the correct CUDA/cuDNN/Python3 modules, export the XLA_FLAGS variable (if needed), and activate the virtualenv by the source command. This includes any time you are submitting to SLURM.