Google ColabでCUDAとtorchのバージョンを変更する。

背景

Google Colaboratory(略称:Colab)では、基本無料でnotebook形式の処理を実行できます。必要なパッケージやGPUでの計算などもできるため簡単に充実した環境を用意できる一方で、インストールされているソフトやパッケージのバージョンがGoogleの意思次第で変わることがあります。
ここでは、意図せずに変わってしまったCUDAとtorchのバージョンに悩まされることなく環境を整備する方法をまとめます。

執筆時の環境と変更先

CUDA:11.0 -> 10.1
torch:1.9.0+cu102 -> 1.8.0

ソースと解説

CUDAとtorchのバージョンを確認します。

!nvcc --version

nvcc: NVIDIA (R) Cuda compiler driver
Copyright (c) 2005-2020 NVIDIA Corporation
Built on Wed_Jul_22_19:09:09_PDT_2020
Cuda compilation tools, release 11.0, V11.0.221
Build cuda_11.0_bu.TC445_37.28845127_0

!python -c 'import torch; print(torch.__version__) '

1.9.0+cu102

変更可能なCUDAのバージョンを確認します。

!ls -d /usr/local/cuda-*
!which nvcc

/usr/local/cuda-10.0 /usr/local/cuda-10.1 /usr/local/cuda-11.0 /usr/local/cuda/bin/nvcc

CUDAのPATHを変更し、変更されたことを確認します。

import os
p = os.getenv('PATH')
ld = os.getenv('LD_LIBRARY_PATH')
os.environ['PATH'] = f"/usr/local/cuda-10.1/bin:{p}"
os.environ['LD_LIBRARY_PATH'] = f"/usr/local/cuda-10.1/lib64:{ld}"
!nvcc --version

nvcc: NVIDIA (R) Cuda compiler driver
Copyright (c) 2005-2019 NVIDIA Corporation
Built on Sun_Jul_28_19:07:16_PDT_2019
Cuda compilation tools, release 10.1, V10.1.243

torchのインストール

インストールには公式の推奨方法を確認してください。
使いたいバージョンによって入れ方が変わる場合があります。
(本記事のように、torchだけでなくtorchvisionやtorchaudioも合わせて入れる場合がある)

!pip install torch==1.8.0 torchvision==0.9.0 torchaudio==0.8.0

Requirement already satisfied: torch==1.8.0 in /usr/local/lib/python3.7/dist-packages (1.8.0)
Requirement already satisfied: torchvision==0.9.0 in /usr/local/lib/python3.7/dist-packages (0.9.0)
Requirement already satisfied: torchaudio==0.8.0 in /usr/local/lib/python3.7/dist-packages (0.8.0)
Requirement already satisfied: numpy in /usr/local/lib/python3.7/dist-packages (from torch==1.8.0) (1.19.5)
Requirement already satisfied: typing-extensions in /usr/local/lib/python3.7/dist-packages (from torch==1.8.0) (3.7.4.3)
Requirement already satisfied: pillow>=4.1.1 in /usr/local/lib/python3.7/dist-packages (from torchvision==0.9.0) (7.1.2)

torchの確認とassert

import torch, torchvision
print(torch.__version__, torch.cuda.is_available())
assert torch.__version__.startswith("1.8")

1.8.0 True

以上で、CUDAとtorchの環境設定が完了し変更後のバージョンで実行するはずです。