Pytorch – 計算を行うデバイスを指定する方法について

概要
Pytorch で計算を指定したデバイス (CPU または GPU) で行う方法について解説します。
Advertisement
torch.device
torch.device は、デバイスを表すクラスです。 torch.Tensor クラスは、計算を実行するデバイスを設定として持っており、計算はそのデバイス上で行われます。テンソルのデバイスは Tensor.device で取得できます。
In [1]:
import torch
a = torch.tensor([1, 2, 3], device=torch.device("cpu"))
b = torch.tensor([1, 2, 3], device=torch.device("cpu"))
c = a * b # a, b の device はどちらも CPU なので、計算は CPU 上で行われる。
print(f"a.device: {a.device}")
print(f"b.device: {b.device}")
print(f"c.device: {c.device}")
a.device: cpu b.device: cpu c.device: cpu
計算を実行する際に、使用するテンソルのデバイスはすべて一致している必要があります。テンソルによってデバイスが異なる場合は、エラーとなります。
In [2]:
import torch
a = torch.tensor([1, 2, 3], device=torch.device("cpu"))
b = torch.tensor([1, 2, 3], device=torch.device("cuda"))
c = a * b # エラー
デバイスの作成方法
- CPU
torch.device("cpu")
- GPU
torch.device("cuda")
- 複数枚の GPU がある場合、使用する GPU の ID を
torch.device("cuda:<GPU ID>")
またはtorch.device("cuda", index=<GPU ID>)
と指定する。 - GPU ID は
nvidia-smi
コマンドで確認できる。1枚目は ID=0、2枚目は ID=1、… となっている。
In [3]:
import torch
# CPU
a = torch.tensor([1, 2, 3], device=torch.device("cpu"))
# GPU
a = torch.tensor([1, 2, 3], device=torch.device("cuda"))
# ID=0 の GPU (GPU が複数枚ある場合)
a = torch.tensor([1, 2, 3], device=torch.device("cuda:0"))
# または
a = torch.tensor([1, 2, 3], device=torch.device("cuda", index=0))
デバイスを引数に指定する場合に、torch.device
オブジェクトの代わりにデバイスを表す文字列を直接指定することもできます。
In [4]:
import torch
# torch.device オブジェクトを指定する場合
a = torch.tensor([1, 2, 3], device=torch.device("cuda:0"))
# 文字列を指定する場合
a = torch.tensor([1, 2, 3], device="cuda:0")
テンソルのデバイスの指定方法
計算を実行するデバイスの指定は、torch.tensor() でテンソルを作成する際に指定する方法と、すでにあるテンソルのデバイスを後から Tensor.to() で変更する方法があります。
- torch.tensor() でテンソルを作成する際に指定する。
- Tensor.to() で後から変更する。
In [5]:
import torch
# torch.tensor() で作成時に device 引数に指定する。
a = torch.tensor([1, 2, 3], device="cuda")
print(a.device) # cpu
# Tensor.to() で後から変更する。
a = torch.Tensor([1, 2, 3])
a = a.to("cuda")
print(a.device) # cuda:0
cuda:0 cuda:0
デバイスを返す関数を作成する
以下のような torch.device
オブジェクトを返す関数を定義しておくと便利かもしれません。
- CUDA が利用可能かどうか
torch.cuda.is_available()
でチェックする。 - GPU を利用する場合は、
gpu_id
で使用する GPU ID を指定する。 - CPU を利用する場合はなにも指定しない。
In [6]:
def get_device(gpu_id=-1):
if gpu_id >= 0 and torch.cuda.is_available():
return torch.device("cuda", gpu_id)
else:
return torch.device("cpu")
device = get_device()
print(device) # cpu
device = get_device(gpu_id=0)
print(device) # cuda:0
cpu cuda:0
Advertisement
デフォルトのデバイスを設定する
テンソルを作成する際にデバイスを指定しない場合、set_default_tensor_type()
で設定されているデバイスが設定されます。
デフォルトでは、CPU になっていますが、以下のように GPU で作成されるように変更できます。
In [7]:
torch.set_default_tensor_type(torch.cuda.FloatTensor)
a = torch.tensor([1, 2, 3])
print(a.device)
cuda:0
-
前の記事
Pytorch – 自作のデータセットを扱う Dataset クラスを作る方法 2020.05.29
-
次の記事
Pytorch – 事前学習モデルを使ってクラス分類モデルを学習する方法 2020.06.01