python: 学習時にGPUを使えるように変更

構成

NVIDIA GeForce RTX 2070 SUPER MSVC 2019 (Microsoft Visual Studio C++) CUDA v10.1 cuDNN 7.4

Python 3.8 TensorFlow-GPU 2.3.0 Keras 2.4.3

手順

（ここから不要かも） MSVC 2019のインストール NVIDIA Drivers for CUDA on WSL のインストール CUDA Toolkit のインストール

sudo apt-key adv --fetch-keys http://developer.download.nvidia.com/compute/cuda/repos/ubuntu1804/x86_64/7fa2af80.pub
sudo sh -c 'echo "deb http://developer.download.nvidia.com/compute/cuda/repos/ubuntu1804/x86_64 /" > /etc/apt/sources.list.d/cuda.list'
sudo apt-get update
sudo apt-get install -y cuda-toolkit-11-0

（ここまで不要かも）

Dockerを使う

NVIDIA Container Toolkit のイントール

#!/bin/sh

distribution=$(. /etc/os-release;echo $ID$VERSION_ID)
curl -s -L https://nvidia.github.io/nvidia-docker/gpgkey | sudo apt-key add -
curl -s -L https://nvidia.github.io/nvidia-docker/$distribution/nvidia-docker.list | sudo tee /etc/apt/sources.list.d/nvidia-docker.list
curl -s -L https://nvidia.github.io/libnvidia-container/experimental/$distribution/libnvidia-container-experimental.list | sudo tee /etc/apt/sources.list.d/libnvidia-container-experimental.list

sudo apt-get update
sudo apt-get install -y nvidia-docker2
sudo usermod -aG docker $USER

build

docker build . -t tensorflow_tomoino --build-arg USER_ID=(id -u) --build-arg GROUP_ID=(id -g) --force-rm=true

コンテナ起動して試してみる

docker run --shm-size=1g --ulimit memlock=-1 --ulimit stack=67108864 --gpus all -v (pwd):/workspace/python --rm -it nvcr.io/nvidia/tensorflow:20.03-tf2-py3

docker run --shm-size=4g --ulimit memlock=-1 --ulimit stack=67108864 --gpus all -v (pwd):/workspace/python --rm -it tensorflow_tomoino:latest

起動後確認

python
>>> import tensorflow as tf
>>> tf.test.gpu_device_name()
>>> print('TensorFlow', tf.__version__)
>>> print('TensorFlow', tf.keras.__version__)

学習

cd /workspace/python/src
python train_cnn.py

モデルの変換

cd /workspace/python
tensorflowjs_converter --input_format keras model/cnn_model.h5 jsmodel

tomoino / Alibi

python: 学習時にGPUを使えるように変更 #31

構成

手順

Dockerを使う

参考