Pytorch with Conda and Cuda
The easiest way to install PyTorch with GPU support is using Conda. This will yield a Python environment, located in your $WORK directory, with a GPU-enabled version of the PyTorch package.
First, you need to open a terminal connection to a node with a GPU. The following command will do this:
srun --partition=2080-galvani --gres=gpu:1 --pty bash
Now you are in a terminal window on a compute node. On the new terminal on the compute node, run the following commands. First, you are creating a new Conda environment with Python version 3.11. Second, you are activating that environment so that you can run commands within it. Third, you are installing the PyTorch package with CUDA/GPU support.
conda create -p $WORK/.conda/py-311-pytorch python=3.11
conda activate $WORK/.conda/py-311-pytorch
conda install pytorch torchvision torchaudio pytorch-cuda=12.1 -c pytorch -c nvidia
Note
Please note that if you are working on the Ferranti Cluster you would need to install pythorch-nightly otherwise when torch is loaded it won't be able to pick up H100 GPUs.
If you need other Python packages in your PyTorch environment, you can install them with python3 -m pip install $PACKAGE_NAME or conda install $PACKAGE_NAME. To find which packages are available, you must search the PyPi website.
If you want to check that this worked, just after running the three commands above, run these two commands at the terminal.
python3 -c "import torch; print(torch.cuda.is_available())"
# Should return True
python3 -c "import torch; print(torch.rand(5, 3))"
# Should return tensor([…]) with 3 columns and 5 rows with random numbers inside.
Created: June 21, 2024