【Pytorch】デバイス(CPU/CUDA)を指定する方法【GPU】
Pytorchでテンソルを扱うデバイスを指定するには、device引数や、.to()メソッドを用いる。
テンソルは以下に定義するものを用いる。
a = torch.ones(1)
print(a)
# 出力
# tensor([1.])GPUでテンソルを扱う
GPUでテンソルを扱うにはテンソルをGPUへ移動する必要がある。
以下のようなコードを書く。
複数の方法があってどれも同じ。
# GPUへの移動(すべて同じ)
b = a.cuda()
print(b)
b = a.to('cuda')
print(b)
b = torch.ones(1, device='cuda')
print(b)
# 出力
# tensor([1.], device='cuda:0')
# tensor([1.], device='cuda:0')
# tensor([1.], device='cuda:0')また、CUDAが使えるかどうかはtorch.cuda.is_available()で分かるので、あらかじめこれでdeviceを取得しておくと便利。
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
print(device)
print(torch.ones(1, device=device))
# 出力
# cuda
# tensor([1.], device='cuda:0')CPUでテンソルを扱う
CPUを指定してテンソルを扱う場合は以下のようなコードを書けばよい。
# CPUへの移動(すべて同じ)
c = b.cpu()
print(c)
c = b.to('cpu')
print(c)
c = torch.ones(1, device='cpu')
print(c)
# 出力
# tensor([1.])
# tensor([1.])
# tensor([1.])