PyTorch学习（2）：torch.device

1. 简述

torch.device 是 PyTorch 中的一个类，它指定了张量（Tensor）和模型应该运行在哪个设备上。在深度学习中，通常会使用 GPU 来加速计算，而 torch.device 允许你轻松地在 CPU 和 GPU 之间切换。

torch.device代表将torch.Tensor分配到的设备的对象。

torch.device包含一个设备类型（'cpu'或'cuda'设备类型）和可选的设备的序号。如果设备序号不存在，则为当前设备;
例如，torch.Tensor用设备构建'cuda'的结果等同于'cuda:X',其中X是torch.cuda.current_device()的结果。

torch.Tensor的设备可以通过Tensor.device访问属性。

在使用torch.device之前，你需要了解当前系统上可用的设备。可以使用 torch.cuda.device_count() 来获取系统上 GPU 的数量。如果系统没有 GPU，这个函数会返回 0。

import torch
''' 获取 GPU 的数量 '''
gpu_count=torch.cuda.device_count()print(f"Number of GPUs available: {gpu_count}")

2. 构造

（1）通过字符串构造设备

torch.device('cpu') 

torch.device('cuda')  # current cuda devicedevice(type='cuda')

（2）通过字符串＋序号构造设备

torch.device('cuda', 0)，device(type='cuda', index=0)

torch.device('cpu', 0)，device(type='cpu', index=0)

（3）可直接使用字符串构建

torch.device函数中的参数通常可以用一个字符串替代。这允许使用代码快速构建原型。

cuda1 = torch.device('cuda:1')

torch.randn((2,3), device=cuda1)

torch.randn((2,3), 'cuda:1')

两种方法等同

（4）有cuda的也可直接输入序号

出于传统原因，可以通过单个设备序号构建设备，将其视为cuda设备。这匹配Tensor.get_device()，它为cuda张量返回一个序数，并且不支持cpu张量。

torch.device(1)

# device(type='cuda', index=1)

（5）创建Tensor时直接指定设备

指定设备的方法可以使用（properly formatted）字符串或（legacy）整数型设备序数，即以下示例均等效：

torch.randn((2,3), device=torch.device('cuda:1'))

torch.randn((2,3), device='cuda:1')

torch.randn((2,3), device=1)  # legacy

3. 使用CPU

可直接使用CPU进行训练，此时torch.device指定CPU。

# 创建一个在 CPU 上运行的设备对象

device = torch.device("cpu")

# 创建一个张量并将其移动到 CPU

tensor_on_cpu = torch.randn(2, 2, device=device)

# 确保张量确实在 CPU 上

print(tensor_on_cpu.device == torch.device("cpu"))

4. 使用GPU

在有GPU的设备上，我们一般使用GPU来加速运算。

# 检查是否有可用的 GPU

if torch.cuda.is_available():

    # 获取第一个 GPU

    device = torch.device("cuda", 0)  # 你可以根据需要更改索引

    # 创建一个张量并将其移动到 GPU

    tensor_on_gpu = torch.randn(2, 2, device=device)

    # 确保张量确实在 GPU 上

    print(tensor_on_gpu.device == torch.device("cuda:0"))

5. 在CPU和GPU之间进行数据的互转

# 创建一个在 CPU 上的张量

tensor_on_cpu = torch.randn(2, 2)

# 将张量移动到 GPU

tensor_on_gpu = tensor_on_cpu.to(device)

# 再将张量移回 CPU

tensor_on_cpu = tensor_on_gpu.to("cpu")

6. 设置默认设备

通常情况下，为了便利，我们可以在工程开头指定默认设备，这样一来，接下来所有的Tensor都会默认使用这个设备。

# 设置默认的 GPU 设备

torch.cuda.set_device(0)


# 现在，所有新创建的张量都会默认在 GPU 上

default_tensor_on_gpu = torch.randn(2, 2)