PyTorchでTensorとモデルのGPU / CPUを指定・切り替え

Posted: | Tags: Python, PyTorch, 機械学習

PyTorchでテンソルtorch.Tensorのデバイス(GPU / CPU)を切り替えるには、to()またはcuda(), cpu()メソッドを使う。torch.Tensorの生成時にデバイス(GPU / CPU)を指定することも可能。

モデル(ネットワーク)すなわちtorch.nn.Moduleのインスタンスにもto()およびcuda(), cpu()メソッドが提供されており、デバイス(GPU / CPU)の切り替えが可能。

ここでは以下の内容について説明する。

  • デバイス(GPU / CPU)を指定してtorch.Tensorを生成
  • torch.Tensorのデバイスを確認: deviceis_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の変更にも用いられる。

dtypedeviceを同時に変更することも可能。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に転送される。第一引数devicetorch.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.Tensorto(), cuda(), cpu()と同じ。

torch.Tensorと異なり、torch.nn.Moduleto(), 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属性がないため、上の例では、便宜上、重みweightdeviceを確認している。

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>)

関連カテゴリー

関連記事