Published on

機械学習まわり環境更新(2021年11月版)

Author
株式会社イエローバックの機械学習エンジニアです

(2021-11-9 追記)

はじめに

機械学習まわりの環境を torch 1.10 など最新のものにしましたので、メモとして残します。 以前の環境との速度比較も少ししました。

環境

  • GPU: RTX2080Ti x 1
  • CPU: Intel Core i9-9900K
  • Memory: 64GB
  • OS: Ubuntu 20.04
  • gcc,g++: 9.3
  • python: 3.9

python インストール

python 3.9 をインストールします。 少し前に python 3.10 がリリースされたのですが、torch 1.10.0 は python3.10 用の precompile version がまだ存在しないため、今回は python のバージョンは 3.9 とします。

$ sudo apt install python3.9 python3.9-venv python3.9-dev

CuDA インストール

pytorch で対応している CuDA は 11.3 なので 11.3 をインストールします。

参照 https://developer.nvidia.com/cuda-downloads?target_os=Linux&target_arch=x86_64&target_distro=Ubuntu&target_version=2004&target_type=debnetwork

レポジトリを設定します。

$ wget https://developer.download.nvidia.com/compute/cuda/repos/ubuntu2004/x86_64/cuda-ubuntu2004.pin
$ sudo mv cuda-ubuntu2004.pin /etc/apt/preferences.d/cuda-repository-pin-600
$ sudo apt-key adv --fetch-keys https://developer.download.nvidia.com/compute/cuda/repos/ubuntu2004/x86_64/7fa2af80.pub
$ sudo add-apt-repository "deb https://developer.download.nvidia.com/compute/cuda/repos/ubuntu2004/x86_64/ /"
$ sudo apt-get update

nVidia ドライバをインストールします。kernel モジュールもインストールされるためインストール後には一度リブートします。

$ sudo apt install nvidia-driver-495 nvidia-dkms-495 libnvidia-gl-495 nvidia-kernel-source-495 nvidia-compute-utils-495 libnvidia-compute-495 libnvidia-common-495 libnvidia-decode-495 libnvidia-encode-495 nvidia-utils-495 xserver-xorg-video-nvidia-495 libnvidia-cfg1-495 libnvidia-fbc1-495
$ sudo reboot

次に CuDA 11.3 をインストールします。

$ sudo apt-get install -y --no-install-recommends cuda-11-3
$ sudo reboot

続いて cuDNN 8.2 をインストールします。

$ sudo apt install libcudnn8=8.2.1.32-1+cuda11.3 libcudnn8-dev=8.2.1.32-1+cuda11.3

TensorRT をインストールします。

$ sudo apt install libnvinfer8=8.0.3-1+cuda11.3 libnvinfer-dev=8.0.3-1+cuda11.3 libnvinfer-plugin8=8.0.3-1+cuda11.3

apt upgrade をしたときに cuda11.5 用などがインストールされないように、バージョンを固定します。

$ sudo apt-mark hold libcudnn8 libcudnn8-dev libnvinfer-dev libnvinfer-plugin8 libnvinfer8

python 仮想環境

python の仮想環境を使用します。

$ python3.9 -m venv ~/.venv/py39-xx
$ source ~/.venv/py39-xx/bin/activate
$ pip install -U pip
$ pip install wheel

pytorch インストール

参照 https://pytorch.org/get-started/locally/

pytorch は cuda 11.3 対応のものをインストールします。

$ pip install \
  torch==1.10.0+cu113 \
  torchvision==0.11.1+cu113 \
  torchaudio==0.10.0+cu113 \
  -f https://download.pytorch.org/whl/cu113/torch_stable.html

apex インストール

参照 https://github.com/NVIDIA/apex

nVidia の apex パッケージをインストールします。

$ git clone https://github.com/NVIDIA/apex
$ cd apex
$ pip install -v --disable-pip-version-check --no-cache-dir --global-option="--cpp_ext" --global-option="--cuda_ext" ./

.
.
Successfully installed apex-0.1

上記メッセージが表示されたら完了です。

jax,flax インストール

参照 https://github.com/google/jax

jax は cuda 11.X, cuDNN 8.2 対応のものをインストールします。

$ pip install jax[cuda11_cudnn82] \
  -f https://storage.googleapis.com/jax-releases/jax_releases.html

続いて flax をインストールします。

$ pip install flax

インストールされたものは以下のバージョンとなりました。

flax              0.3.6
jax               0.2.24
jaxlib            0.1.73+cuda11.cudnn82

transformers などをインストールします。

$ pip install transformers fugashi unidic-lite ipadic datasets deepspeed

インストール状況確認

$ ds_report
--------------------------------------------------
DeepSpeed C++/CUDA extension op report
--------------------------------------------------
NOTE: Ops not installed will be just-in-time (JIT) compiled at
      runtime if needed. Op compatibility means that your system
      meet the required dependencies to JIT install the op.
--------------------------------------------------
JIT compiled ops requires ninja
ninja .................. [OKAY]
--------------------------------------------------
op name ................ installed .. compatible
--------------------------------------------------
cpu_adam ............... [NO] ....... [OKAY]
cpu_adagrad ............ [NO] ....... [OKAY]
fused_adam ............. [NO] ....... [OKAY]
fused_lamb ............. [NO] ....... [OKAY]
sparse_attn ............ [NO] ....... [OKAY]
transformer ............ [NO] ....... [OKAY]
stochastic_transformer . [NO] ....... [OKAY]
async_io ............... [NO] ....... [OKAY]
transformer_inference .. [NO] ....... [OKAY]
utils .................. [NO] ....... [OKAY]
quantizer .............. [NO] ....... [OKAY]
--------------------------------------------------
DeepSpeed general environment info:
torch install path ............... ['/home/ono/.venv/py39-xx/lib/python3.9/site-packages/torch']
torch version .................... 1.10.0+cu113
torch cuda version ............... 11.3
nvcc version ..................... 11.3
deepspeed install path ........... ['/home/ono/.venv/py39-xx/lib/python3.9/site-packages/deepspeed']
deepspeed info ................... 0.5.5, unknown, unknown
deepspeed wheel compiled w. ...... torch 1.10, cuda 11.3

また主要パッケージ類は以下の通りです。

datasets           1.15.1
deepspeed          0.5.5
jax                0.2.24
jaxlib             0.1.73+cuda11.cudnn82
torch              1.10.0+cu113
transformers       4.12.2
tokenizers         0.10.3

全体は、以下の gist に保存しました。 https://gist.github.com/yellowback/30f4a17828192dd8fd45aa6d90e0af11

速度比較

以前のメイン環境と学習速度、評価速度を pytorch 版で比較してみました。 livedoor コーパス分類の学習および評価を実行し、1 秒あたりの学習サンプル数(train_samples_per_second)と 1 秒あたりの評価サンプル数(eval_samples_per_second)を比較しました。 値が大きいほど速度が早いことを意味します。

なお、モデルやパラメタは以下の通りです。

  • model: cl-tohoku/bert-base-japanese-whole-word-masking
  • batchsize: 8
  • epoch: 3
  • deepspeed: 未使用
  • fp16_opt_level: O2
  • fp16_backend: amp/apex

以前の環境

  • python: 3.8
  • CuDA: 11.1
  • pytorch: 1.8.1+cu111
  • transformers: 4.10.3

今回の環境

  • python: 3.9
  • CuDA: 11.3
  • pytorch: 1.10.0+cu113
  • transformers: 4.12.2

速度

fp16_backend = amp

環境学習サンプル数/秒評価サンプル数/秒
以前の環境40.082.6
今回の環境44.3160.5

fp16_backend = apex

環境学習サンプル数/秒評価サンプル数/秒
以前の環境59.0219.8
今回の環境58.0210.6

apex 使用時には、少しだけ性能が落ちています。 amp 使用時にはとくに評価時の性能が向上しています。

おわりに

機械学習環境を更新しました。

  • Ubuntu: 20.04
  • python: 3.9
  • CuDA: 11.3
  • cuDNN: 8.2.1
  • pytorch: 1.10.0
  • jax: 0.2.24
  • jaxlib: 0.1.73
  • flax: 0.3.6
  • transformers: 4.12.2
Discuss on Twitter