Research IT

PyTorch

PyTorch is a premier open-source machine learning and deep learning framework.  It has quickly become of the most prominent frameworks available. There are two ways to install PyTorch 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 torch-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/torch-mamba

And install PyTorch (Pip is recommended inside Micromamba for PyTorch to ensure you get the exact CUDA binaries):

pip install torch torchvision torchaudio

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/torch-venv

Then we activate the environment:

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

And last we install PyTorch using pip:

pip install torch torchvision torchaudio

Testing the Environment

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

import torch
import time

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
print(f"Using device: {device}")
if torch.cuda.is_available():
   print(f"GPU Name: {torch.cuda.get_device_name(0)}")

print("Creating matrices...")
size = 5000
x = torch.randn(size, size, device=device)
y = torch.randn(size, size, device=device)

print("Running matrix multiplication...")
start_time = time.time()
z = torch.matmul(x, y)

if torch.cuda.is_available():
   torch.cuda.synchronize()
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_torch.py and run it:

python ./test_torch.py

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

Using device: cuda
GPU Name: NVIDIA A100-SXM4-80GB
Creating matrices...
Running matrix multiplication...
Multiplication complete! Shape: torch.Size([5000, 5000])
Time taken: 0.0619 seconds

Tutorials

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