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

ops/2024/12/20 20:12:31/

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

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

  1. 检查 PyTorch 版本
    确保使用的是 PyTorch 2.0 或更高版本:

    python">import torch
    print(torch.__version__)  # 必须是 2.0 或更高版本
    
  2. 检查 CUDA 是否可用
    在设置 GPU 为默认设备之前,确认 CUDA 可用性:

    python">print(torch.cuda.is_available())  # True 表示可用
    
  3. 设置默认设备为 GPU
    使用 torch.set_default_device 将默认设备设置为 GPU:

    python">import torch# 确保 CUDA 可用
    if torch.cuda.is_available():# 设置默认设备为 GPUtorch.set_default_device('cuda')print("默认设备已设置为 GPU")
    else:print("CUDA 不可用,无法设置 GPU 为默认设备")
    
  4. 验证默认设备设置
    创建一个张量,验证其是否在 GPU 上:

    python">x = torch.tensor([1.0, 2.0, 3.0])
    print(x.device)  # 输出:cuda:0
    
  5. 模型自动加载到 GPU
    如果设置了默认设备,模型的参数和新建的张量会自动加载到 GPU:

    python">class MyModel(torch.nn.Module):def __init__(self):super().__init__()self.linear = torch.nn.Linear(10, 1)def forward(self, x):return self.linear(x)model = MyModel()
    print(next(model.parameters()).device)  # 输出:cuda:0
    
全局设置代码示例

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

python">import torch# 检查并设置默认设备
if torch.cuda.is_available():torch.set_default_device('cuda')print("默认设备已设置为 GPU")
else:raise RuntimeError("CUDA 不可用,请检查环境配置")# 示例:自动使用 GPU 创建张量
x = torch.tensor([1.0, 2.0, 3.0])
print(f"x device: {x.device}")  # 输出:cuda:0# 示例:自动将模型参数放到 GPU
model = torch.nn.Linear(5, 2)
print(f"Model parameters device: {next(model.parameters()).device}")  # 输出:cuda:0
注意事项
  1. 与设备显式管理的代码兼容性
    如果代码中显式指定了设备(如 tensor.to(device)),torch.set_default_device 不会影响这些张量。建议在全局设置后,尽量减少显式设备管理操作。

  2. 多 GPU 环境
    如果有多个 GPU,可以指定具体设备,比如 'cuda:1'。示例:

    python">torch.set_default_device('cuda:1')  # 使用第二块 GPU
    
  3. 性能调优
    默认将所有操作转移到 GPU 可能并不适合所有场景,尤其是小规模任务时,GPU 的初始化开销可能超过性能提升。根据需求灵活调整设备。


http://www.ppmy.cn/ops/143567.html

相关文章

JavaScript九宫格随机抽奖示例

<!DOCTYPE html> <html lang"en"><head><meta charset"utf-8"><title>九宫格抽奖</title><style>/* 全局样式重置 */* {margin: 0;padding: 0;}/* 抽奖容器样式 */.lottery-container {width: 300px;height: 30…

面试题整理5----进程、线程、协程区别及僵尸进程处理

面试题整理5----进程、线程、协程区别及僵尸进程处理 1. 进程、线程与协程的区别1.1 进程&#xff08;Process&#xff09;1.2 线程&#xff08;Thread&#xff09;1.3 协程&#xff08;Coroutine&#xff09;2. 总结对比 3. 僵尸进程3.1 什么是僵尸进程&#xff1f;3.2 僵尸进…

数据可视化-1. 折线图

目录 1. 折线图适用场景分析 1. 1 时间序列数据展示 1.2 趋势分析 1.3 多变量比较 1.4 数据异常检测 1.5 简洁易读的数据可视化 1.6 特定领域的应用 2. 折线图局限性 3. 折线图代码实现 3.1 Python 源代码 3.2 折线图效果&#xff08;网页显示&#xff09; 1. 折线图…

【Java基础面试题020】Java中Exception和Error有什么区别?

回答重点 Exception和Error都是Throwable类的子类&#xff08;在Java代码中只有继承了Throwable类的实例&#xff0c;才可以被throw和被catch&#xff09;他们表示程序运行时发生的异常或错误情况 总结来看&#xff0c;Exception表示可以被处理的程序异常&#xff0c;Error表…

力扣-图论-14【算法学习day.64】

前言 ###我做这类文章一个重要的目的还是给正在学习的大家提供方向和记录学习过程&#xff08;例如想要掌握基础用法&#xff0c;该刷哪些题&#xff1f;&#xff09;我的解析也不会做的非常详细&#xff0c;只会提供思路和一些关键点&#xff0c;力扣上的大佬们的题解质量是非…

The Rise and Potential of Large Language ModelBased Agents:A Survey---讨论

讨论 论法学硕士研究与Agent研究的互利性 近年来&#xff0c;随着激光诱导金属化技术的发展&#xff0c;激光诱导金属化与化学剂交叉领域的研究取得了长足的进步&#xff0c;促进了这两个领域的发展。在此&#xff0c;我们期待着LLM研究和Agent研究相互提供的一些益处和发展机…

D98【python 接口自动化学习】- pytest进阶之fixture用法

day98 pytest的fixture功能之session 学习日期&#xff1a;20241215 学习目标&#xff1a;pytest基础用法 -- pytest的fixture功能之session 学习笔记&#xff1a; fixture(scop"session") (scop"session") 是多个文件调用一次&#xff0c;.py文件就是…

【JS/TS鼠标气泡跟随】文本提示 / 操作提示

适用于任何类型项目&#xff1a;vue、react、angular、js、ts、jsp、jquery 1、功能封装&#xff1a; export function useMouseActionTip(text: string, parentEl: HTMLElement, offset?: XY) {function mousemove(e: MouseEvent) {const offsetX offset?.x || 16;const of…