Skip to content

PyTorch

PyTorch is an optimized tensor library for deep learning using GPUs and CPUs.

Homepage Version(s): 2.0

Installing PyTorch on the HPC#

PyTorch is not currently in the standard Anaconda package provided on the HPC, but you can install it yourself using conda.

First, create a conda environment (full documentation is available on our Conda page):

Warning

Inititial installation may take a while to run and use approximately 5.5G in the directory/folder in which it is installed.

1
2
3
4
5
6
7
8
# Create conda environment (1)
(base) [USER@h22-login-26 ~]$ conda create -n pytorch_app

# Activate conda environment
(base) [USER@h22-login-26 ~]$ conda activate pytorch_app

# Install PyTorch (2) 
(pytorch_app) [USER@h22-login-26 ~]$ conda install pytorch torchvision torchaudio pytorch-cuda=11.7 -c pytorch -c nvidia
  1. If you haven't run conda init yet, please refer to our instructions for how to do so
  2. This command assumes you want to use Pytorch with GPUs. To customize the installation command, please refer to the getting started page on the official docs

Using PyTorch on the HPC#

Use the sample code from the Verification section of the PyTorch guide:

1
2
3
import torch
x = torch.rand(5, 3)
print(x)

Using GPUs with PyTorch#

To run a PyTorch job on a GPU node, you need to submit a job to our job scheduler, and indicate that you want to run using GPUs.

Warning

Running torch.cuda.is_available() on the login node will return False. It will return True on a GPU node.

Example Slurm GPU Submit script#

1
2
3
4
5
6
7
8
9
#!/bin/bash

#SBATCH -A backfill2   # Submit to backfill2 (note, this account has a four-hour walltime limit)
#SBATCH -n 12          # Request resources for 12 tasks
#SBATCH --gres=gpu:1   # Request a single GPU card

conda activate pytorch_app  # Substitute your conda environment for 'pytorch_app'

srun python MYCODE.py