PyTorch 2.0 中设置默认使用 GPU 的方法

PyTorch 2.0 中设置默认使用 GPU 的方法

在 PyTorch 2.0 中，默认环境下仍旧是使用 CPU 进行盘算，除非明白指定使用 GPU。torch.set_default_device 是 PyTorch 2.0 引入的新功能，用于设置默认设备，使得全部后续张量和模块在没有明白指定设备的环境下，会被创建在这个默认设备上。这在代码中提供了一种更简便的方式来指定设备，而无需在每次创建张量或模型时手动指定。

• 查抄 PyTorch 版本
  确保使用的是 PyTorch 2.0 或更高版本：
  
  1. import torch
  2. print(torch.__version__)  # 必须是 2.0 或更高版本

  复制代码
• 查抄 CUDA 是否可用
  在设置 GPU 为默认设备之前，确认 CUDA 可用性：
  
  1. print(torch.cuda.is_available())  # True 表示可用

  复制代码
• 设置默认设备为 GPU
  使用 torch.set_default_device 将默认设备设置为 GPU：
  
  1. import torch
  3. # 确保 CUDA 可用
  4. if torch.cuda.is_available():
  5.     # 设置默认设备为 GPU
  6.     torch.set_default_device('cuda')
  7.     print("默认设备已设置为 GPU")
  8. else:
  9.     print("CUDA 不可用，无法设置 GPU 为默认设备")

  复制代码
• 验证默认设备设置
  创建一个张量，验证其是否在 GPU 上：
  
  1. x = torch.tensor([1.0, 2.0, 3.0])
  2. print(x.device)  # 输出：cuda:0

  复制代码
• 模型自动加载到 GPU
  如果设置了默认设备，模型的参数和新建的张量会自动加载到 GPU：
  
  1. class MyModel(torch.nn.Module):
  2.     def __init__(self):
  3.         super().__init__()
  4.         self.linear = torch.nn.Linear(10, 1)
  6.     def forward(self, x):
  7.         return self.linear(x)
  9. model = MyModel()
  10. print(next(model.parameters()).device)  # 输出：cuda:0

  复制代码

全局设置代码示例

以下代码展示如何在脚本中全局设置默认设备为 GPU：

1. import torch
3. # 检查并设置默认设备
4. if torch.cuda.is_available():
5.     torch.set_default_device('cuda')
6.     print("默认设备已设置为 GPU")
7. else:
8.     raise RuntimeError("CUDA 不可用，请检查环境配置")
10. # 示例：自动使用 GPU 创建张量
11. x = torch.tensor([1.0, 2.0, 3.0])
12. print(f"x device: {x.device}")  # 输出：cuda:0
14. # 示例：自动将模型参数放到 GPU
15. model = torch.nn.Linear(5, 2)
16. print(f"Model parameters device: {next(model.parameters()).device}")  # 输出：cuda:0

复制代码

注意事项

• 与设备显式管理的代码兼容性
  如果代码中显式指定了设备（如 tensor.to(device)），torch.set_default_device 不会影响这些张量。建议在全局设置后，只管镌汰显式设备管理操作。
  
• 多 GPU 环境
  如果有多个 GPU，可以指定具体设备，比如 'cuda:1'。示例：
  
  1. torch.set_default_device('cuda:1')  # 使用第二块 GPU

  复制代码
• 性能调优
  默认将全部操作转移到 GPU 可能并不得当全部场景，尤其是小规模任务时，GPU 的初始化开销可能超过性能提升。根据需求机动调解设备。
  

免责声明：如果侵犯了您的权益，请联系站长，我们会及时删除侵权内容，谢谢合作！更多信息从访问主页：qidao123.com:ToB企服之家，中国第一个企服评测及商务社交产业平台。