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

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

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  # エラー

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

import torch

# torch.device オブジェクトを指定する場合
a = torch.tensor([1, 2, 3], device=torch.device("cuda:0"))

# 文字列を指定する場合
a = torch.tensor([1, 2, 3], device="cuda:0")

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

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

torch.set_default_tensor_type(torch.cuda.FloatTensor)
a = torch.tensor([1, 2, 3])
print(a.device)