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)