Using PyTorch on Berzelius

1. Introduction
2. Installing PyTorch via Conda/Mamba
3. Installing PyTorch via pip
4. Installing PyTorch via Apptainer

1. Introduction

PyTorch is an open-source machine learning library for Python that provides a flexible and efficient framework for building and training deep neural networks. PyTorch is commonly used for various machine learning tasks, including deep learning, natural language processing (NLP), computer vision, reinforcement learning, and more.

On Berzelius, PyTorch can be installed using various methods, such as Conda/Mamba or pip, following the official installation instructions. Alternatively, you can also install and run PyTorch within an Apptainer container for improved portability and reproducibility.

2. Installing PyTorch via Conda/Mamba

It’s considered good practice to specify the version of the main package to install—in this case, PyTorch—to ensure compatibility and reproducibility.

module load Miniforge3/24.7.1-2-hpc1-bdist
mamba create --name pytorch-2.5.0-python-3.10 python=3.10
mamba activate pytorch-2.5.0-python-3.10
mamba install pytorch==2.5.0 torchvision torchaudio pytorch-cuda=12.1 -c pytorch -c nvidia

To check if PyTorch detects the GPU:

python -c "import torch; print('GPU available: ' + str(torch.cuda.is_available()))"

Installation instructions for previous versions of PyTorch can be found here.

The official PyTorch channel on Anaconda (-c pytorch) is being deprecated and is no longer actively maintained. PyTorch 2.5.1 will be the last release published to the pytorch channel. Going forward, users are encouraged to install PyTorch exclusively from the conda-forge channel.

module load Miniforge3/24.7.1-2-hpc1-bdist
mamba create --name pytorch-2.5.1-python-3.10 python=3.10
mamba activate pytorch-2.5.1-python-3.10
CONDA_OVERRIDE_CUDA=11.8 mamba install pytorch==2.5.1=cuda* torchvision=*=cuda* torchaudio=*=cuda*

3. Installing PyTorch via pip

module load Miniforge3/24.7.1-2-hpc1-bdist
mamba create --name pytorch-2.5.0-python-3.10 python=3.10
mamba activate pytorch-2.5.0-python-3.10
pip3 install torch==2.5.0 torchvision torchaudio --index-url https://download.pytorch.org/whl/cu121

4. Installing PyTorch via Apptainer

We can build an Apptainer image using the following definition file pytorch_2.0.1.def. To learn more refer to the Apptainer User Guide.

Bootstrap: docker
From: nvidia/cuda:11.7.1-cudnn8-devel-ubuntu22.04

%environment

export PATH=/opt/mambaforge/bin:$PATH
export PYTHONNOUSERSITE=True

%post

apt-get update && apt-get install -y --no-install-recommends \
git \
nano \
wget \
curl

# Install Mambaforge
cd /tmp
curl -L -O "https://github.com/conda-forge/miniforge/releases/latest/download/Mambaforge-$(uname)-$(uname -m).sh"
bash Mambaforge-$(uname)-$(uname -m).sh -fp /opt/mambaforge -b
rm Mambaforge*sh

export PATH=/opt/mambaforge/bin:$PATH

mamba install python=3.10 pytorch==2.0.1 torchvision torchaudio torchdata torchtext pytorch-cuda=11.7 -c pytorch -c nvidia -y

# Pin packages
cat <<EOT > /opt/mambaforge/conda-meta/pinned
pytorch==2.0.1
EOT

mamba install matplotlib jupyterlab -y

We build the image from the definition file:

apptainer build pytorch_2.0.1 pytorch_2.0.1.def

The Apptainer image can be easily extended with more packages and software by modifying the definition file and rebuilding the image.

Using PyTorch on Berzelius

1. Introduction

2. Installing PyTorch via Conda/Mamba

3. Installing PyTorch via pip

4. Installing PyTorch via Apptainer

User support

Getting access

Everything OK!

Self-service