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

相关推荐
沐雪架构师12 分钟前
AI大模型开发原理篇-2:语言模型雏形之词袋模型
人工智能·语言模型·自然语言处理
python算法(魔法师版)1 小时前
深度学习深度解析:从基础到前沿
人工智能·深度学习
小王子10241 小时前
设计模式Python版 组合模式
python·设计模式·组合模式
kakaZhui1 小时前
【llm对话系统】大模型源码分析之 LLaMA 位置编码 RoPE
人工智能·深度学习·chatgpt·aigc·llama
struggle20252 小时前
一个开源 GenBI AI 本地代理(确保本地数据安全),使数据驱动型团队能够与其数据进行互动,生成文本到 SQL、图表、电子表格、报告和 BI
人工智能·深度学习·目标检测·语言模型·自然语言处理·数据挖掘·集成学习
佛州小李哥2 小时前
通过亚马逊云科技Bedrock打造自定义AI智能体Agent(上)
人工智能·科技·ai·语言模型·云计算·aws·亚马逊云科技
Mason Lin3 小时前
2025年1月22日(网络编程 udp)
网络·python·udp
清弦墨客3 小时前
【蓝桥杯】43697.机器人塔
python·蓝桥杯·程序算法
云空3 小时前
《DeepSeek 网页/API 性能异常(DeepSeek Web/API Degraded Performance):网络安全日志》
运维·人工智能·web安全·网络安全·开源·网络攻击模型·安全威胁分析
AIGC大时代3 小时前
对比DeepSeek、ChatGPT和Kimi的学术写作关键词提取能力
论文阅读·人工智能·chatgpt·数据分析·prompt