🔥

PyTorch：GPUの指定方法

に公開
Python
PyTorch
CUDA
mps
tech

PyTorch：GPUの指定方法

TL;DR

以下の3通りの方法でGPUを指定できる。

  1. to()メソッドでオブジェクト毎に指定
  2. with torch.device()でスコープ毎に指定
  3. torch.set_default_device()で全体に指定

1. to()メソッドでオブジェクト毎に指定

import torch


x = torch.tensor([1, 2, 3]).to("cuda:0")
print(x.device)
# > cuda:0

2. with torch.device()でスコープ毎に指定

import torch


with torch.device("cuda:0"):
    x = torch.tensor([1, 2, 3])
    y = torch.tensor([4, 5, 6])
    print(x.device, y.device)
    # > cuda:0, cuda:0

注意

with torch.device()のスコープ内でデバイスを指定する以前に作成されたテンソルはデバイスが変更されので、
異なるデバイス同士の計算を行うと普通にエラーが発生する。

import torch


x = torch.tensor([1, 2, 3])
with torch.device("cuda:0"):
    y = torch.tensor([4, 5, 6])
    print(x.device, y.device)
    # > cpu, cuda:0
    z = x + y
    # > RuntimeError: Expected all tensors to be on the same device, but found at least two devices, cuda:0 and cpu!

3. torch.set_default_device()で全体に指定

import torch


torch.set_default_device("cuda:0")
x = torch.tensor([1, 2, 3])
y = torch.tensor([4, 5, 6])
print(x.device, y.device)
# > cuda:0, cuda:0

まとめ

set_default_deviceを使用すればcuda <-> mpsなどで実行環境が違ってもコードの修正箇所が少なくて済むので便利。
matplotlibのようなcpu処理を行う時のみwith torch.device()を使うと綺麗に書ける。

参考

Discussion