PyTorchでTensorとモデルのGPU / CPUを指定・切り替え
PyTorchでテンソルtorch.Tensor
のデバイス(GPU / CPU)を切り替えるには、to()
またはcuda()
, cpu()
メソッドを使う。torch.Tensor
の生成時にデバイス(GPU / CPU)を指定することも可能。
- torch.Tensor.to() — PyTorch 1.7.1 documentation
- torch.Tensor.cuda() — PyTorch 1.7.1 documentation
- torch.Tensor.cpu() — PyTorch 1.7.1 documentation
モデル(ネットワーク)すなわちtorch.nn.Module
のインスタンスにもto()
およびcuda()
, cpu()
メソッドが提供されており、デバイス(GPU / CPU)の切り替えが可能。
- torch.nn.Module.to() — PyTorch 1.7.1 documentation
- torch.nn.Module.cuda() — PyTorch 1.7.1 documentation
- torch.nn.Module.cpu() — PyTorch 1.7.1 documentation
ここでは以下の内容について説明する。
- デバイス(GPU / CPU)を指定して
torch.Tensor
を生成 torch.Tensor
のデバイスを確認:device
、is_cuda
torch.Tensor
のGPU / CPUを切り替え:to()
,cuda()
,cpu()
- 環境に応じてGPU / CPUを切り替える方法
- マルチGPU環境におけるコンテキストマネージャ
torch.cuda.device
- モデル(ネットワーク)のGPU / CPU切り替え
なお、当然ながら、GPUが使えない環境でCPUからGPUに転送するといった処理を行うとエラーになるので注意。
本記事のサンプルコードにおけるPyTorchのバージョンは以下の通り。バージョンが異なると仕様が異なる場合もあるので注意。
import torch
print(torch.__version__)
# 1.7.1
デバイス(GPU / CPU)を指定してtorch.Tensorを生成
torch.tensor()
やtorch.ones()
, torch.zeros()
などのtorch.Tensor
を生成する関数では、引数device
を指定できる。
以下のサンプルコードはtorch.tensor()
だが、torch.ones()
などでも同じ。
引数device
にはtorch.device
のほか、文字列をそのまま指定することもできる。
CPUの場合。デフォルトはCPU。
print(torch.tensor([0.1, 0.2]))
# tensor([0.1000, 0.2000])
print(torch.tensor([0.1, 0.2], device=torch.device('cpu')))
# tensor([0.1000, 0.2000])
print(torch.tensor([0.1, 0.2], device='cpu'))
# tensor([0.1000, 0.2000])
GPUの場合。cuda:n
のように、n
にGPUの番号(インデックス)を指定する。省略した場合はデフォルトのGPUが使われる。GPU番号をそのまま整数値で指定することもできる。
print(torch.tensor([0.1, 0.2], device=torch.device('cuda:0')))
# tensor([0.1000, 0.2000], device='cuda:0')
print(torch.tensor([0.1, 0.2], device=torch.device('cuda')))
# tensor([0.1000, 0.2000], device='cuda:0')
print(torch.tensor([0.1, 0.2], device=torch.device(0)))
# tensor([0.1000, 0.2000], device='cuda:0')
print(torch.tensor([0.1, 0.2], device='cuda:0'))
# tensor([0.1000, 0.2000], device='cuda:0')
print(torch.tensor([0.1, 0.2], device='cuda'))
# tensor([0.1000, 0.2000], device='cuda:0')
print(torch.tensor([0.1, 0.2], device=0))
# tensor([0.1000, 0.2000], device='cuda:0')
使用できるGPUの数などの確認については以下の記事を参照。
torch.Tensorのデバイスを確認: device、is_cuda
torch.Tensor
のデバイスはdevice
属性で取得できる。
t_cpu = torch.tensor([0.1, 0.2])
print(t_cpu.device)
# cpu
print(type(t_cpu.device))
# <class 'torch.device'>
t_gpu = torch.tensor([0.1, 0.2], device='cuda')
print(t_gpu.device)
# cuda:0
print(type(t_gpu.device))
# <class 'torch.device'>
GPUかどうかはis_cuda
属性で確認できる。PyTorch1.7.1
時点ではis_cpu
属性は提供されていない。
print(t_cpu.is_cuda)
# False
print(t_gpu.is_cuda)
# True
torch.TensorのGPU / CPUを切り替え: to(), cuda(), cpu()
既存のtorch.Tensor
のデバイス(GPU / CPU)を切り替える(転送する)にはto()
, cuda()
, cpu()
メソッドを使う。
to()
の場合。引数device
を指定する。torch.tensor()
などと同様に、引数device
にはtorch.device
、文字列、GPUの場合は数値を指定できる。
t_cpu = torch.tensor([0.1, 0.2])
print(t_cpu.device)
# cpu
t_gpu = t_cpu.to('cuda')
print(t_gpu.device)
# cuda:0
to()
はデータ型dtype
の変更にも用いられる。
dtype
とdevice
を同時に変更することも可能。to(device, dtype)
の順番だと位置引数として指定できるが、to(dtype, device)
の順番だとキーワード引数として指定する必要があるので注意。
print(t_cpu.to('cuda', torch.float64))
# tensor([0.1000, 0.2000], device='cuda:0', dtype=torch.float64)
# print(t_cpu.to(torch.float64, 'cuda'))
# TypeError: to() received an invalid combination of arguments - got (torch.dtype, str), but expected one of:
# * (torch.device device, torch.dtype dtype, bool non_blocking, bool copy, *, torch.memory_format memory_format)
# * (torch.dtype dtype, bool non_blocking, bool copy, *, torch.memory_format memory_format)
# * (Tensor tensor, bool non_blocking, bool copy, *, torch.memory_format memory_format)
print(t_cpu.to(dtype=torch.float64, device='cuda'))
# tensor([0.1000, 0.2000], device='cuda:0', dtype=torch.float64)
cuda()
の場合。引数を省略した場合はデフォルトのGPUに転送される。第一引数device
にtorch.device
、文字列、数値を指定して、任意のGPUを指定可能。存在しないGPU番号やCPUを指定するとエラー。
print(t_cpu.cuda().device)
# cuda:0
print(t_cpu.cuda(0).device)
# cuda:0
# print(t_cpu.cuda(1).device)
# RuntimeError: CUDA error: invalid device ordinal
print(t_cpu.cuda('cuda:0').device)
# cuda:0
# print(t_cpu.cuda('cpu').device)
# RuntimeError: Invalid device, must be cuda device
cpu()
の場合。
print(t_gpu.cpu().device)
# cpu
to()
, cuda()
, cpu()
メソッドのいずれも、既に指定したデバイスに保存されている場合は元のオブジェクトをそのまま返す。新たにコピーが生成されてしまうようなことはない。
t_cpu2 = t_cpu.to('cpu')
print(t_cpu is t_cpu2)
# True
t_gpu2 = t_gpu.cuda()
print(t_gpu is t_gpu2)
# True
複数のGPUが使用できる環境においてはコンテキストマネージャtorch.cuda.device
を用いて対象のGPUを指定する方法もある。後述。
環境に応じてGPU / CPUを切り替える方法
GPUが使用可能な環境かどうかはtorch.cuda.is_available()
で判定できる。
GPUが使える環境ではGPUを、そうでない環境でCPUを使うようにするには、例えば以下のように適当な変数(ここではdevice
)にtorch.device
を代入しておいて、引数device
に指定すればよい。三項演算子を利用している。
device = torch.device('cuda:0') if torch.cuda.is_available() else torch.device('cpu')
print(device)
# cuda:0
t = torch.tensor([0.1, 0.2], device=device)
print(t.device)
# cuda:0
マルチGPU環境におけるコンテキストマネージャtorch.cuda.device
複数のGPUを使用できる環境においてはコンテキストマネージャtorch.cuda.device
を用いて、ブロック内でのデフォルトGPUを指定できる。
使い方は以下の公式ドキュメントのサンプルコードを参照。
モデル(ネットワーク)のGPU / CPU切り替え
PyTorchにおけるモデル(ネットワーク)はすべてtorch.nn.Module
を継承している。
torch.nn.Module
のメソッドとして、to()
, cuda()
, cpu()
が提供されており、モデルのデバイス(GPU / CPU)を切り替えられる。
torch.nn.Linear
を例とする。便宜上、torch.manual_seed()
でシードを固定している。
torch.manual_seed(0)
net = torch.nn.Linear(2, 2)
print(isinstance(net, torch.nn.Module))
# True
to()
, cuda()
, cpu()
でGPU / CPUを切り替える。デバイスの指定方法はtorch.Tensor
のto()
, cuda()
, cpu()
と同じ。
torch.Tensor
と異なり、torch.nn.Module
のto()
, cuda()
, cpu()
はin-placeの処理。呼び出し元のオブジェクト自体が更新される。
print(net.weight.device)
# cpu
net.to('cuda')
print(net.weight.device)
# cuda:0
net.cpu()
print(net.weight.device)
# cpu
net.cuda()
print(net.weight.device)
# cuda:0
なお、torch.nn.Module
にはdevice
属性がないため、上の例では、便宜上、重みweight
のdevice
を確認している。
Modules can hold parameters of different types on different devices, and so it’s not always possible to unambiguously determine the device. https://github.com/pytorch/pytorch/issues/7460
モデルとテンソルが異なるデバイスに存在していると、モデルにテンソルを流して実行する際にエラーになるので注意。
t_gpu = torch.tensor([0.1, 0.2], device='cuda')
print(t_gpu.device)
# cuda:0
print(net(t_gpu))
# tensor([-0.1970, 0.0273], device='cuda:0', grad_fn=<AddBackward0>)
t_cpu = torch.tensor([0.1, 0.2], device='cpu')
print(t_cpu.device)
# cpu
# print(net(t_cpu))
# RuntimeError: Tensor for 'out' is on CPU, Tensor for argument #1 'self' is on CPU, but expected them to be on GPU (while checking arguments for addmm)
上述の通り、環境に応じてGPU / CPUを切り替えるにはtorch.cuda.is_available()
で判定すればよい。
device = torch.device('cuda:0') if torch.cuda.is_available() else torch.device('cpu')
t = torch.tensor([0.1, 0.2], device=device)
torch.manual_seed(0)
net = torch.nn.Linear(2, 2)
net.to(device)
print(net(t_gpu))
# tensor([-0.1970, 0.0273], device='cuda:0', grad_fn=<AddBackward0>)