DEV Community

Cover image for Upgrade to PyTorch 2.0
Zoheb Abai
Zoheb Abai

Posted on

Upgrade to PyTorch 2.0

Why Upgrade?

Upgrade Objectives

  • Python ≥ 3.8, ≤ 3.11
  • CUDA ≥ 11.7.0
  • CUDNN ≥ 8.5.0.96
  • Pytorch ≥ 2.0.0

“We expect that with PyTorch 2, people will change the way they use PyTorch day-to-day”
“Data scientists will be able to do with PyTorch 2.x the same things that they did with 1.x, but they can do them faster and at a larger scale”
— Soumith Chintala


Steps for Upgrade

If you have Python ≥ 3.8, ≤ 3.11 jump to next section

Steps for upgrading Python from ≤ 3.8 to 3.10

For Clean Installation remove all existing Python related files

# Replace X with the specific version number
sudo apt --purge remove python3.X
sudo apt-get autoremove
sudo apt-get autoclean
Enter fullscreen mode Exit fullscreen mode

Pre-Installations Actions

sudo apt update

# Install required dependencies
sudo apt install build-essential zlib1g-dev libncurses5-dev libgdbm-dev libnss3-dev libssl-dev libreadline-dev libffi-dev wget
Enter fullscreen mode Exit fullscreen mode

Installing Python 3.10.6 from source

Download desired version (here 3.10.6) from Python Website

# Extract the source code
tar -xvf Python-3.10.6.tgz

# Configure the build 
cd python-3.10.6
./configure --enable-optimizations --prefix=/usr/local

# Start the build process
make -j $(nproc)

# Once the build completes, install Python
sudo make install
Enter fullscreen mode Exit fullscreen mode

Open ./bashrcfile and add following lines at the end

export PATH="/usr/local/bin:$PATH"

Save the file and update the environment variables for the current session by running

source ~/.bashrc

Verify Python version

python3 --version

which python3
Enter fullscreen mode Exit fullscreen mode

If you already have CUDA ≥ 11.7.0 jump to next section

Steps for upgrading Cuda ≤ 11.7 on Ubuntu 22.04 with a Nvidia Geforce RTX Graphics Card:

For Clean Installation remove all existing cuda related files

sudo apt-get --purge remove "*cuda*" "*cublas*" "*cufft*" "*cufile*" "*curand*"  "*cusolver*" "*cusparse*" "*gds-tools*" "*npp*" "*nvjpeg*" "nsight*"
sudo apt-get --purge remove "*nvidia*"
sudo apt-get autoremove
sudo apt-get autoclean
Enter fullscreen mode Exit fullscreen mode

Pre-Installation Actions

# Verify You Have a CUDA-Capable GPU
lspci | grep -i nvidia

# Verify the System Has gcc Installed
gcc --version

# Verify the System has the Correct Kernel Headers and Development Packages Installed
sudo apt-get install linux-headers-$(uname -r)
Enter fullscreen mode Exit fullscreen mode

Install NVIDIA CUDA Toolkit 11.7.1 (Debian Installer Preferred)

# Install the repository meta-data, update the GPG key, update the apt-get cache, and install CUDA:
wget https://developer.download.nvidia.com/compute/cuda/repos/ubuntu2204/x86_64/cuda-ubuntu2204.pin
sudo mv cuda-ubuntu2204.pin /etc/apt/preferences.d/cuda-repository-pin-600
wget https://developer.download.nvidia.com/compute/cuda/11.7.1/local_installers/cuda-repo-ubuntu2204-11-7-local_11.7.1-515.65.01-1_amd64.deb
sudo dpkg -i cuda-repo-ubuntu2204-11-7-local_11.7.1-515.65.01-1_amd64.deb
sudo cp /var/cuda-repo-ubuntu2204-11-7-local/cuda-*-keyring.gpg /usr/share/keyrings/
sudo apt-get update
sudo apt-get -y install cuda
Enter fullscreen mode Exit fullscreen mode

For more details, check this.

During Cuda Installation you might get asked for creating a password for MOK management — Do it.

Reboot the system to load the NVIDIA drivers. If you get a blue screen, DO NOT continue to boot, instead enrol the key providing the password you created a while ago. And then continue to boot.

Open ./bashrc file and add following lines at the end

export PATH=/usr/local/cuda-11.7/bin${PATH:+:${PATH}}
export LD_LIBRARY_PATH=/usr/local/cuda-11.7/lib64${LD_LIBRARY_PATH:+:${LD_LIBRARY_PATH}}
Enter fullscreen mode Exit fullscreen mode

Save the file and update the environment variables for the current session by running

source ~/.bashrc

Verify the Cuda version

nvcc --version

nvidia-smi
Enter fullscreen mode Exit fullscreen mode

If you already have CUDNN ≥ 8.5.0.96 jump to next section

Steps for upgrading CUDNN ≤ 8.5.0.96

Install CUDNN 8.5.0.96 (Debian Installer Preferred)

wget https://developer.nvidia.com/compute/cudnn/secure/8.5.0/local_installers/11.7/cudnn-local-repo-ubuntu2204-8.5.0.96_1.0-1_amd64.deb
sudo dpkg -i cudnn-local-repo-ubuntu2204-8.5.0.96_1.0-1_amd64.deb

# Import the CUDA GPG key
sudo cp /var/cudnn-local-repo-ubuntu2204-8.5.0.96/cudnn-local-*-keyring.gpg /usr/share/keyrings/

# Refresh the repository metadata
sudo apt-get update

# Install the runtime library
sudo apt-get install libcudnn8=8.5.0.96-1+cuda11.7

# Install the developer library
sudo apt-get install libcudnn8-dev=8.5.0.96-1+cuda11.7
Enter fullscreen mode Exit fullscreen mode

Fore more details, check this.


If you already have PyTorch ≥ 2.0.0 you are awesome.

Steps for upgrading PyTorch≤ 2.0.0

# If you have virtualenv and use pip as manager
python3 -m pip install torch==2.0.0+cu117 torchvision==0.15.1+cu117 torchaudio==2.0.1 --index-url https://download.pytorch.org/whl/cu117
Enter fullscreen mode Exit fullscreen mode

For other OS or package manager, check this.

For downloading wheel files, check this.

Verify PyTorch 2.0 Installation

python3 -c "import torch; print(torch.__version__)"
Enter fullscreen mode Exit fullscreen mode

If you get errors google it or comment it down.

I hope the article helped. Thanks.


Top comments (0)