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

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():
        # 设置默认设备为 GPU
        torch.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 的初始化开销可能超过性能提升。根据需求灵活调整设备。

相关推荐
学术小八几秒前
2025年人工智能、虚拟现实与交互设计国际学术会议
人工智能·交互·vr
nbsaas-boot1 小时前
Java 正则表达式白皮书:语法详解、工程实践与常用表达式库
开发语言·python·mysql
仗剑_走天涯1 小时前
基于pytorch.nn模块实现线性模型
人工智能·pytorch·python·深度学习
chao_7891 小时前
二分查找篇——搜索旋转排序数组【LeetCode】两次二分查找
开发语言·数据结构·python·算法·leetcode
cnbestec2 小时前
协作机器人UR7e与UR12e:轻量化设计与高负载能力助力“小而美”智造升级
人工智能·机器人·协作机器人·ur协作机器人·ur7e·ur12e
zskj_zhyl2 小时前
毫米波雷达守护银发安全:七彩喜跌倒检测仪重构居家养老防线
人工智能·安全·重构
gaosushexiangji3 小时前
利用sCMOS科学相机测量激光散射强度
大数据·人工智能·数码相机·计算机视觉
ai小鬼头4 小时前
AIStarter新版重磅来袭!永久订阅限时福利抢先看
人工智能·开源·github
说私域5 小时前
从品牌附庸到自我表达:定制开发开源AI智能名片S2B2C商城小程序赋能下的营销变革
人工智能·小程序
飞哥数智坊5 小时前
新版定价不够用,Cursor如何退回旧版定价
人工智能·cursor