torch.device函数

torch.device 是 PyTorch 中用于表示计算设备(如CPU或GPU)的类。它允许你在代码中指定你希望在哪个设备上执行张量和模型操作,本文主要介绍了 torch.device 函数的用法和功能。

1.创建设备对象

通过 torch.device(“device_type:device_id”) 来创建一个设备对象,其中 device_type 是设备类型,可以是 “cpu” 或 “cuda”,而 device_id 是GPU的标识号(从 0 开始)

cpu_device = torch.device("cpu")          # 创建一个CPU设备对象
gpu_device = torch.device("cuda:0")       # 创建第一个GPU设备对象

2.将张量和模型移动到设备:

一旦你创建了设备对象,你可以使用 .to(device) 方法将张量或模型移动到特定的设备上。这对于在不同设备之间进行数据传递和计算操作非常有用。

tensor_cpu = torch.tensor([1, 2, 3])       # 创建一个CPU张量
tensor_cpu = tensor_cpu.to(cpu_device)     # 将张量移动到CPU上

model = MyModel()                         # 创建一个模型
model = model.to(gpu_device)              # 将模型移动到GPU上

3.检查设备类型:

你可以使用 .type 属性来检查一个设备对象的类型。

print(cpu_device.type)      # 输出: 'cpu'
print(gpu_device.type)      # 输出: 'cuda'

4.切换设备:

在代码中,你可以随时切换使用不同的设备,这对于在有GPU和没有GPU的环境中编写代码非常有用

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

5.devices = torch.device(“cuda”)

devices = torch.device("cuda")  # 创建一个表示GPU设备的对象
devices = torch.device("cpu")  # 创建一个表示CPU设备的对象

这样,你就可以将模型和其他操作移到GPU(或CPU)上进行加速运算。然后,你可以将模型和其他需要在GPU(或CPU)上执行的张量发送到 devices 中,使它们在GPU(或CPU)上进行计算。

6.devices = torch.device(“cuda:%s”%(args.gpu) if torch.cuda.is_available() else “cpu”)

这行代码是一个常见的用法,它根据是否有可用的GPU来动态选择设备。

# 检查是否有可用的GPU,如果有则选择使用GPU,否则使用CPU
devices = torch.device("cuda:%s" % (args.gpu) if torch.cuda.is_available() else "cpu")

# 这行代码会执行以下步骤:
# 1. torch.cuda.is_available() 检查是否有可用的GPU。
# 2. 如果有可用的GPU,则将 "cuda:<gpu_id>" 分配给 devices,其中 <gpu_id> 是 args.gpu 变量中指定的GPU设备号。
# 3. 如果没有可用的GPU,则将 "cpu" 分配给 devices,表示使用CPU设备。

这段代码的目的是根据系统中是否有可用的GPU来选择合适的设备。它首先检查是否存在可用的GPU(通过 torch.cuda.is_available()),如果存在,就将设备设置为指定的GPU(“cuda:<gpu_id>”),否则将设备设置为CPU(“cpu”)。这种方法使代码能够在有GPU的情况下充分利用GPU进行加速,而在没有GPU的情况下仍然可以在CPU上运行。

cuda:%s 是一个格式化字符串,用于在字符串中插入变量值。%s 是一个占位符,表示在这个位置插入一个变量的值。

% (args.gpu): 这是格式化字符串的实际值,将会替换掉 %s。args.gpu 是一个变量,它可能是一个GPU设备号,将在这里插入。

综合起来,cuda:%s 表示一个字符串模板,其中的 %s 将被替换为某个值,这个值取决于 args.gpu 变量的实际值。

例如,如果 args.gpu 的值为 “0”,那么 “cuda:%s” % (args.gpu) 将被替换为 “cuda:0”,表示使用 GPU 0。如果 args.gpu 的值为 “1”,那么 “cuda:%s” % (args.gpu) 将被替换为 “cuda:1”,表示使用 GPU 1,以此类推(本段落均为英文引号)。

devices = torch.device(“cuda:%s”%(args.gpu) if torch.cuda.is_available() else “cpu”)和devices = torch.device(“cuda”)这两个代码片段都是用于选择设备(GPU或CPU)的方式,但在一些方面存在一些细微的差异。
devices = torch.device(“cuda:%s”%(args.gpu) if torch.cuda.is_available() else “cpu”)

这个代码片段首先检查是否有可用的GPU(通过 torch.cuda.is_available() 函数)。
如果有可用的GPU,则它会尝试将设备设置为指定的GPU,即 “cuda:<gpu_id>”,其中 <gpu_id> 是 args.gpu 中指定的GPU设备号。
如果没有可用的GPU,它将设备设置为CPU。
这种方式在你想要根据情况切换设备(GPU或CPU)时很有用。例如,你可能希望在有GPU时使用GPU进行训练,而在没有GPU时回退到CPU。

devices = torch.device(“cuda”)

这个代码片段直接将设备设置为默认的第一个可用GPU(通常是GPU 0)。
如果没有可用的GPU,将引发一个错误。
这种方式适用于你确定你要在特定的GPU上运行,并且你希望代码在没有GPU时报错。

如果你希望根据是否有可用的GPU来自动选择设备并在没有GPU时使用CPU,可以使用第一个代码片段。

如果你确定你要使用特定的GPU并且希望在没有GPU时引发错误,可以使用第二个代码片段。

总之,这两种方式都是为了在代码中选择合适的设备,以便在GPU上加速计算。选择使用哪种方式取决于你的项目需求和硬件配置。

猜你喜欢

转载自blog.csdn.net/thy0000/article/details/132073049
今日推荐