Zenn
📺

自前のGPUでPyTorchを動かすための前準備(Windows)

2021/08/31に公開
Docker
WSL
PyTorch
GPU
CUDA
tech

はじめに

念願のGPU(RTX3070)を積んだPC(Windows)を購入したので、PyTorchを動かしてディープラーニングすることにしました。
色々と前準備が必要だったので、メモとして残しました。

手順

  1. CUDA on WSLの導入
  2. PyTorchのインストール

1. CUDA on WLSの導入

これについては、素晴らしい記事があるのでこちらを参照すると良いです。
こちらの記事の CUDA Toolkit のインストール までを行えばOKです。

インストールするパッケージについては、公式のページを参考にしました。

2. PyTorchのインストール

PyTorchのインストールですが、ローカルのAnacondaを使う方法とdockerを使う方法の両方を試しました。
お好きな方を参照してください。

ローカルのAnacondaでのインスール

Anacondaはすでにwsl上に導入されているものとします。
CUDA11.1に対応しているPyTorchは、以下のコマンドでインストールできます。

pip install torch==1.9.0+cu111 torchvision==0.10.0+cu111 torchaudio===0.9.0 -f https://download.pytorch.org/whl/torch_stable.html

どのバージョンのPyTorchをインストールすればよいのかについては、こちらで確認できます。

dockerを使用したインストール

dockerを使用する場合は、nvidiaが用意しているPyTorchのイメージを使用する方法と、pythonイメージを使用してPyTorchをインストールする方法があります。
前者の方法だと、追加でライブラリをインストールしたい場合、condaとpipのコンフリクトが起きるケースがあったので、ここでは後者の方法で実施します。

以下のような、pythonイメージのDockerfileを用意します。

Dockerfile
FROM python:3

# 必要なライブラリのインストール
ADD requirements.txt .
RUN pip install --upgrade pip && \
    pip install --no-cache-dir --default-timeout=1000  -r requirements.txt

# PyTorchのインストール
RUN pip install --no-cache-dir --default-timeout=1000 \
    torch==1.9.0+cu111 \
    torchvision==0.10.0+cu111 \
    -f https://download.pytorch.org/whl/torch_stable.html

Dockerfileを使用してビルドします。

docker build -t {image_name} .

image_nameには好きな名前を入れて下さい。
このイメージを以下のコマンドで起動します。

docker run \
    --gpus all \
    -it \
    --rm \
    -v {local_dir}:{container_dir} \
    --ipc=host \
    {image_name}

local_dir、container_dirは絶対パスで指定して下さい。

GPUが使用できるか確認

PyTorchの公式ページにサンプルのコードがあったので、動作するか確認してみます。

# -*- coding: utf-8 -*-

import torch
import math


dtype = torch.float
device = torch.device("cuda:0")
print(device)

# Create random input and output data
x = torch.linspace(-math.pi, math.pi, 2000, device=device, dtype=dtype)
y = torch.sin(x)

# Randomly initialize weights
a = torch.randn((), device=device, dtype=dtype)
b = torch.randn((), device=device, dtype=dtype)
c = torch.randn((), device=device, dtype=dtype)
d = torch.randn((), device=device, dtype=dtype)

learning_rate = 1e-6
for t in range(2000):
    # Forward pass: compute predicted y
    y_pred = a + b * x + c * x ** 2 + d * x ** 3

    # Compute and print loss
    loss = (y_pred - y).pow(2).sum().item()
    if t % 100 == 99:
        print(t, loss)

    # Backprop to compute gradients of a, b, c, d with respect to loss
    grad_y_pred = 2.0 * (y_pred - y)
    grad_a = grad_y_pred.sum()
    grad_b = (grad_y_pred * x).sum()
    grad_c = (grad_y_pred * x ** 2).sum()
    grad_d = (grad_y_pred * x ** 3).sum()

    # Update weights using gradient descent
    a -= learning_rate * grad_a
    b -= learning_rate * grad_b
    c -= learning_rate * grad_c
    d -= learning_rate * grad_d


print(f'Result: y = {a.item()} + {b.item()} x + {c.item()} x^2 + {d.item()} x^3')

下図のように問題なく動いていれば、導入完了です。
実行結果

Discussion