PyTorch 2.0 以下版本中设置默认使用 GPU 的方法

PyTorch 2.0 以下版本中设置默认使用 GPU 的方法

在 PyTorch 2.0以下版本中,默认情况下仍然是使用 CPU 进行计算,除非明确指定使用 GPU。在 PyTorch 2.0 以下版本中,虽然没有 torch.set_default_device 的便捷方法,但可以通过显式地将张量、模型和操作分配到 GPU 来使用 GPU。

1. 检查 GPU 可用性

在使用 GPU 之前,首先检查系统中是否有可用的 GPU。

python 复制代码
import torch

# 检查是否有可用的 GPU
print(torch.cuda.is_available())  # 返回 True 或 False

# 检查可用 GPU 的数量
print(torch.cuda.device_count())

# 当前 GPU 名称
if torch.cuda.is_available():
    print(torch.cuda.get_device_name(0))

2. 将张量移动到 GPU

张量可以通过 .to('cuda').cuda() 方法显式地移动到 GPU。

python 复制代码
# 创建一个张量并将其移动到 GPU
x = torch.tensor([1.0, 2.0, 3.0])
x_gpu = x.to('cuda')  # 或 x.cuda()
print(x_gpu.device)  # 输出:cuda:0

# 在 GPU 上进行计算
y = x_gpu * 2
print(y)  # 输出在 GPU 上的结果
3. 将模型移动到 GPU

PyTorch 模型及其参数需要显式地移动到 GPU。

python 复制代码
# 定义一个简单的模型
model = torch.nn.Linear(10, 1)

# 将模型移动到 GPU
model = model.to('cuda')  # 或 model.cuda()

# 检查模型参数所在的设备
print(next(model.parameters()).device)  # 输出:cuda:0
4. 确保输入数据和模型在同一设备上

模型和输入数据需要在同一个设备上,否则会报错。

python 复制代码
# 创建一个张量并移动到 GPU
input_data = torch.randn(5, 10).to('cuda')

# 定义并移动模型到 GPU
model = torch.nn.Linear(10, 1).to('cuda')

# 前向传播
output = model(input_data)
print(output)

5. 使用 torch.device 动态管理设备

可以使用 torch.device 动态管理设备。

python 复制代码
# 定义设备
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')

# 将张量移动到设备
x = torch.tensor([1.0, 2.0, 3.0]).to(device)

# 将模型移动到设备
model = torch.nn.Linear(10, 1).to(device)

6. 优化器和损失函数的设备兼容性

当使用 GPU 时,模型的输出和目标(target)都需要在同一设备上。

python 复制代码
# 创建数据和目标,并移动到 GPU
data = torch.randn(5, 10).to('cuda')
target = torch.randn(5, 1).to('cuda')

# 定义模型并移动到 GPU
model = torch.nn.Linear(10, 1).to('cuda')

# 定义损失函数
criterion = torch.nn.MSELoss()

# 定义优化器
optimizer = torch.optim.SGD(model.parameters(), lr=0.01)

# 前向传播
output = model(data)
loss = criterion(output, target)

# 反向传播
loss.backward()
optimizer.step()

7. 混合设备计算(可选)

在多 GPU 或混合 CPU/GPU 环境中,可以手动管理每个张量或模型的设备。

python 复制代码
# 在 CPU 上创建张量
x_cpu = torch.tensor([1.0, 2.0, 3.0])

# 在 GPU 上创建张量
x_gpu = x_cpu.to('cuda')

# 将结果移动回 CPU
result = x_gpu * 2
result_cpu = result.to('cpu')
print(result_cpu)

总结

在 PyTorch 2.0 以下版本中,使用 GPU 的核心是 显式地将张量和模型移动到 GPU,并确保所有相关操作在同一设备上完成。以下是核心方法的汇总:

  • 检查 GPU 可用性: torch.cuda.is_available()

  • 移动张量到 GPU: .to('cuda').cuda()

  • 移动模型到 GPU: .to('cuda').cuda()

  • 动态设备管理: torch.device

相关推荐
米小虾2 小时前
AI 开始改进 AI:OpenAI 公布自我改进时间表,同一天自家首席科学家喊停
人工智能·openai
故七月2 小时前
告别 AI 时代品牌 “隐身”:万域智瞰 AI‑GEO,构建品牌大模型时代营销新基建
大数据·人工智能
lifallen2 小时前
短暂 Agent 与持久化工程:让理解脱离对话
人工智能·学习·ai·重构·ai编程
知识分享小能手2 小时前
深度学习学习教程,从入门到精通,深度模型中的优化 — 完整知识点与代码案例(8)
人工智能·深度学习·学习
javaDocker2 小时前
金融级 AIOps 架构跃迁(从“1-5-10“快恢目标到 AI Agent 告警收敛的完整技术实践)
人工智能·金融·架构
青 春 记 忆2 小时前
零基础入门python65:FastAPI 安全调用大模型API
python·fastapi·后端开发
多看书少吃饭3 小时前
GPT‑6 Astra 与 AGI 时代:当 AI 开始承担完整的工作
人工智能·gpt·agi
浩风祭月3 小时前
GPT-6 Astra API怎么接?Responses API迁移、异步工具调用与兼容项清单
人工智能·chatgpt·大模型·openai api·gpt-6 astra
FPC工厂——皇榜科技3 小时前
FPC拼板的利用率怎么算?排版方式不同,材料成本差15%
人工智能·科技·pcb工艺
Raspberry_Pi_官方账号3 小时前
Raspberry Pi+TinyML:基于边缘计算的低成本离线皮肤癌AI诊断系统
人工智能·边缘计算·树莓派·raspberry pi·tinyml