PyTorchでのGPUの使用方法

1. 概要

今回はPyTorchのGPUの使用方法を紹介したいと思います。

PyTorchだとGPUを使用する際には、いろいろと指定する必要があります。

GPU環境がない場合は、Google Colaboratoryなどを使って試してみてください。

2. 基本的な使い方

まずは自身の環境でGPUが使えるかを確認します。

import torch
print(torch.cuda.is_available())
# Outut
True

Trueと出力されていればGPUが使えます。

次にGPUのデバイス名を変数に格納します。

device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')
print(device)
# Output
cuda:0

バイス名を変数に格納したら、計算対象の変数をGPUに送り、

あとは計算させるだけになります。

# 変数を用意
a = torch.randn(5)
b = torch.randn(5)

# 変数をGPUに送る
a = a.to(device)
b = b.to(device)

print(a + b)
# Output
# 乱数を使っているので実行するたびに結果が変わります。
tensor([ 0.8977,  1.4709, -1.2835,  0.8451,  2.7441], device='cuda:0')

GPUを使う際の注意点として、計算する変数同士はGPUに送らた状態で計算されなければなりません。

片方をGPUに送っただけだと、エラーが出てしまいます。

# 変数を用意
a = torch.randn(5)
b = torch.randn(5)

# 今回は変数aだけをGPUに送る。
a = a.to(device)

print(a + b)
# Output
---------------------------------------------------------------------------
RuntimeError                              Traceback (most recent call last)
<ipython-input-12-e81e582b6fa9> in <module>()
----> 1 a + b

RuntimeError: Expected all tensors to be on the same device, but found at least two devices, cuda:0 and cpu!