TensorFlow is an end-to-end open source platform for machine learning. It has a comprehensive, flexible ecosystem of tools, libraries and community resources that lets researchers push the state-of-the-art in ML and developers easily build and deploy ML powered applications.

This document will cover how to install TensorFlow (TF) by using Anaconda and run a toy example in Python. Make sure review Anaconda and setup the recommended configuration first, otherwise, you will fill out your home directory by Conda environments.

TensorFlow Environment

We have two options for using TensorFlow (TF), one is using CPUs for the computational tasks and the other one is using GPU resources. Depends on your workflow and identity of your calculations, you might choose any of these methods.

To use TF with CPU resources, use the following to create the environment (env) and install TF:

srun -p Lewis --mem 16G -c 2 --pty /bin/bash
module load miniconda3
conda create --name tensorflow-env tensorflow

And use the following to install TF for GPUs:

srun -p Lewis --mem 16G -c 2 --pty /bin/bash
module load miniconda3
conda create --name tensorflow-gpu-env tensorflow-gpu

Note that tensorflow-gpu includes cudatoolkit and cudnn in addition to TF.

A Simple Example

After creating the preferred env, we can use Python to run TF. The following is a simple example from TensorFlow tutorials. This example uses the MNIST database and Keras to run a pattern recognition training task on images of handwritten digits.

Create a Python file called minist.py including the following scripts:

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(input_shape=(28, 28)),
  tf.keras.layers.Dense(128, activation='relu'),
  tf.keras.layers.Dense(10, activation='softmax')


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

To run the above code with CPU, we can create a batch file called tf-cpu.sh:


#SBATCH -p Lewis
#SBATCH --mem 16G
#SBATCH -c 4

module load miniconda3
conda activate tensorflow-env

python3 minist.py

And submit it by sbatch tf-cpu.sh. You can find the Slurm output file when job is done.

Also, if you have created tensorflow-gpu-env environment, you can use the following batch file, called tf-gpu.sh, to run this job with a GPU.


#SBATCH -p Gpu
#SBATCH --gres gpu:1
#SBATCH --mem 16G
#SBATCH -c 4

module load miniconda3
conda activate tensorflow-gpu-env

python3 minist.py

Now use sbatch tf-gpu.sh to run the job. Slurm out file will show the outputs.

To learn more about submitting and monitoring jobs, review out training page in here. And learn more about using GPU resources in our how to page.