Comment on page

Installing Frameworks (PyTorch, TensorFlow, Jax)

This page describes installing popular frameworks like TensorFlow, PyTorch & JAX, etc. on your Oscar account.
Preface: Oscar is a heterogeneous cluster meaning we have nodes with different architecture GPUs (Pascal, Volta, Turing, and Ampere). We recommend building the environment first time on Ampere GPUs with the latest CUDA11 modules so it's backward compatible with older architecture GPUs.
In this example, we will install PyTorch (refer to sub-pages for TensorFlow and Jax).
Step 1: Request an interactive session on a GPU node with Ampere architecture GPUs
interact -q gpu -g 1 -f ampere -m 20g -n 4
Here, -f = feature. We only need to build on Ampere once.
Step 2: Once your session has started on a compute node, run nvidia-smi to verify the GPU and then load the appropriate modules
module load python/3.11.0 openssl/3.0.0
Step 3: Create and activate the virtual environment
python3 -m venv pytorch.venv
source pytorch.venv/bin/activate
Step 4: Install the required packages
pip install --upgrade pip
pip install torch torchvision torchaudio -f
The aforementioned will install the latest version of PyTorch with cuda11 compatibility, for older versions you can specify the version by:
pip install torch==1.8.0+cu111 torchvision==0.9.0+cu111 torchaudio==0.8.0 -f
Step 5: Test that PyTorch is able to detect GPUs
>>> import torch
>>> torch.cuda.get_device_name(0)
'NVIDIA GeForce RTX 3090'
If the above functions return True and GPU model, then it's working correctly. You are all set, now you can install other necessary packages.