1. 概要
今回はPyTorchのGPUの使用方法を紹介したいと思います。
PyTorchだとGPUを使用する際には、いろいろと指定する必要があります。
GPU環境がない場合は、Google Colaboratoryなどを使って試してみてください。
2. 基本的な使い方
まずは自身の環境でGPUが使えるかを確認します。
import torch print(torch.cuda.is_available())
# Outut True
Trueと出力されていれば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!
リンク