PytorchでのGPUの使用方法

環境の確認

このページではすでにNVIDIA driverとCUDA libralyがinstallされ、環境構築が終了しているものとする。

GPUの確認

以下のコマンドを、Macならterminal, Windowsならコマンドプロンプトに入力実行することでGPUの状態を確認できる。

nvidia-smi

出力の例は以下の通り。（4GPUの例）

Thu Apr  2 00:21:04 2020       
+-----------------------------------------------------------------------------+
| NVIDIA-SMI 430.14       Driver Version: 430.14       CUDA Version: 10.2     |
|-------------------------------+----------------------+----------------------+
| GPU  Name        Persistence-M| Bus-Id        Disp.A | Volatile Uncorr. ECC |
| Fan  Temp  Perf  Pwr:Usage/Cap|         Memory-Usage | GPU-Util  Compute M. |
|===============================+======================+======================|
|   0  GeForce RTX 208...  On   | 00000000:01:00.0 Off |                  N/A |
| 41%   30C    P8    16W / 260W |      0MiB / 11019MiB |      0%      Default |
+-------------------------------+----------------------+----------------------+
|   1  GeForce RTX 208...  On   | 00000000:02:00.0 Off |                  N/A |
| 41%   35C    P0    59W / 260W |      0MiB / 11019MiB |      0%      Default |
+-------------------------------+----------------------+----------------------+
|   2  GeForce RTX 208...  On   | 00000000:03:00.0 Off |                  N/A |
| 41%   29C    P8    10W / 260W |      0MiB / 11019MiB |      0%      Default |
+-------------------------------+----------------------+----------------------+
|   3  GeForce RTX 208...  On   | 00000000:07:00.0 Off |                  N/A |
| 41%   28C    P8    14W / 260W |      0MiB / 11019MiB |      0%      Default |
+-------------------------------+----------------------+----------------------+
                                                                               
+-----------------------------------------------------------------------------+
| Processes:                                                       GPU Memory |
|  GPU       PID   Type   Process name                             Usage      |
|=============================================================================|
|  No running processes found                                                 |
+-----------------------------------------------------------------------------+

上の出力の様に、使用可能なGPUの分だけGPU情報が表示される。例だと4GPUが認識されていることがわかる。Memory Useageの列を見ると0MiB / 11019MiB と表示され、GPUのメモリが現在使われていないことがわかる。GPUの使用中にnvidia-smiを使用すると500MiB/11010MiBのようにどのくらいメモリを専有しているのかがわかる。会社や研究室などで共有のGPUを使用しているときは、使用状況を確認する必要がある。

CUDAの確認方法

python上でtorch.cuda.is_availableの返り値を見ることで、CUDAが現在利用可能かどうかを知ることができる。

import torch
print(torch.cuda.is_available())

#出力:Ture
  • CUDAが利用可能 => True
  • CUDAが利用できない状態 => False

GPUを使用する

GPUにTensorを載せる(1)

一番簡単な方法は、Tensorを初期化するときにdeviceを指定することである。

import tensor
X = tensor.randn(3, 2, device="cuda:0")

cuda:0を指定することで、上の例で言うXをGPUに転送することができる。cuda:1のように数字を変えると指定したGPUに転送される。

以上でGPUの使用方法がわかったが、実際には以下のように記述することが多い。

GPUにTensorを載せる(2)

すべてのPytorchのTensorにはメンバ関数のto()が含まれている。このto()は、どのデバイスにTensorを置くかを指定することができる。デバイスの初期化はtorch.deviceオブジェクトで行う。

  • CPUを使うとき：torch.deviceに"cpu"を入力して初期化
  • GPUを使うとき：torch.deviceに"cuda:0″を入力して初期化
    • 複数のGPUが存在するときはコロンのあとの番号でGPU番号を指定。

ここまでの内容を踏まえ、PytorchでGPUを使用する場合の一般的な流れを示す。

  1. GPUが使えるかどうか確認
  2. GPUが使えるなら"cuda:0″,そうでないなら"cpu"でtorch.deviceを初期化
  3. Tensorを初期化(初期化段階ではTensorはcpu上にある)
  4. Tensorを2. で指定したデバイスに載せる

コードは次のようになる。

import torch
#1
if torch.cuda.is_available():
    d_type = "cuda:0"
else:
    d_type = "cpu"
#2
device = torch.device(d_type)
#3 適当なTensorを用意
A = torch.randn(6, 3)
#4
A.to(device)

実際は次のようにまとめて記述する事が多い。

import torch
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
A = torch.randn(3, 5)
A.to(device)

Tensorが使っているGPUを確認する

GPUに置かれたTensorにはget_device()という関数が用意されている。これを使うと、どのGPUに置かれているのか知ることができる。

>>>A.get_device()
0

ちなみにcpuに置かれたTensorで同じことをするとエラーになる。
違うGPUに乗っているTensor同士で計算を行うとエラーが発生するので注意。

pytorchdeeplearning,gpu,python,pytorch,人工知能,深層学習

Posted by kazu