在PyTorch中使用GPU加速:从基础操作到模型部署

本文将通过具体代码示例,详细介绍如何在PyTorch中利用GPU进行张量计算和模型训练,包含设备查询、数据迁移以及模型部署等完整流程。


1. 查看GPU硬件信息

使用 nvidia-smi 命令检查GPU状态和进程信息:

python 复制代码
# 查看GPU信息
!nvidia-smi

输出示例

显示当前GPU型号(如NVIDIA GeForce RTX 3060)、显存使用情况和运行中的进程。


2. 计算设备管理

2.1 定义设备对象

python 复制代码
import torch
from torch import nn

# 定义不同设备对象
device_cpu = torch.device('cpu')
device_gpu0 = torch.cuda.device('cuda')
device_gpu1 = torch.cuda.device('cuda:1')
device_cpu, device_gpu0, device_gpu1

输出

bash 复制代码
(device(type='cpu'), <torch.cuda.device at 0x...>, <torch.cuda.device at 0x...>)

2.2 查询可用GPU数量

python 复制代码
torch.cuda.device_count()  # 返回可用GPU数量

输出1(表示系统中有1块可用GPU)

3. 灵活选择设备

3.1 按需选择GPU或CPU

python 复制代码
def try_gpu(i=0):
    """选择指定GPU,若不可用则返回CPU"""
    if torch.cuda.device_count() >= i + 1:
        return torch.device(f'cuda:{i}')
    return torch.device('cpu')

def try_all_gpu():
    """返回所有可用GPU,若无则返回CPU"""
    devices = [torch.device(f'cuda:{i}') for i in range(torch.cuda.device_count())]
    return devices if devices else [torch.device('cpu')]

try_gpu(), try_gpu(10), try_all_gpu()

输出

bash 复制代码
(device(type='cuda', index=0), device(type='cpu'), [device(type='cuda', index=0)])

4. 张量与GPU操作

4.1 默认设备查询

python 复制代码
x = torch.tensor([1, 2, 3])
x.device  # 默认在CPU上

输出device(type='cpu')

4.2 显存中创建张量

python 复制代码
x_gpu = torch.ones(2, 3, device='cuda:0')  # 在GPU 0上创建全1张量
x_gpu

输出

bash 复制代码
tensor([[1., 1., 1.],
        [1., 1., 1.]], device='cuda:0')

4.3 GPU间计算

python 复制代码
y_gpu = torch.rand(2, 3, device='cuda:0')
x_gpu + y_gpu  # 必须同一设备才能计算

输出示例

bash 复制代码
tensor([[1.4571, 1.4172, 1.1364],
        [1.7572, 1.3977, 1.2402]], device='cuda:0')

5. 神经网络模型部署到GPU

5.1 模型迁移

python 复制代码
# 定义一个简单神经网络
net = nn.Sequential(nn.Linear(3, 1))
net = net.to(device='cuda')  # 将模型参数移至GPU

# 输入数据需与模型在同一设备
output = net(x_gpu)
output

输出示例

bash 复制代码
tensor([[-0.4271],
        [-0.4271]], device='cuda:0', grad_fn=<AddmmBackward0>)

5.2 验证参数位置

python 复制代码
net[0].weight.data.device  # 检查参数存储位置

输出device(type='cuda', index=0)

6. 关键注意事项

  1. 设备一致性:输入数据和模型必须在同一设备(CPU/GPU)上才能进行计算。

  2. 显存管理 :及时释放不再使用的GPU张量(del tensor + torch.cuda.empty_cache())。

  3. 多GPU支持 :可通过torch.nn.DataParallel实现多卡并行训练。


总结

本文演示了PyTorch中GPU加速的核心操作,包括设备选择、张量迁移和模型部署。合理利用GPU可显著加速深度学习任务,建议在训练大型模型时优先使用GPU环境。如果遇到CUDA相关错误,请检查驱动版本和PyTorch的GPU支持安装。

相关推荐
qinyia41 分钟前
Wisdom SSH 是一款创新性工具,通过集成 AI 助手,为服务器性能优化带来极大便利。
服务器·人工智能·ssh
昨日之日20063 小时前
Wan2.2-S2V - 音频驱动图像生成电影级质量的数字人视频 ComfyUI工作流 支持50系显卡 一键整合包下载
人工智能·音视频
一个天蝎座 白勺 程序猿5 小时前
Python爬虫(47)Python异步爬虫与K8S弹性伸缩:构建百万级并发数据采集引擎
爬虫·python·kubernetes
XiaoMu_0016 小时前
基于Django+Vue3+YOLO的智能气象检测系统
python·yolo·django
SEO_juper6 小时前
大型语言模型SEO(LLM SEO)完全手册:驾驭搜索新范式
人工智能·语言模型·自然语言处理·chatgpt·llm·seo·数字营销
攻城狮7号6 小时前
腾讯混元翻译模型Hunyuan-MT-7B开源,先前拿了30个冠军
人工智能·hunyuan-mt-7b·腾讯混元翻译模型·30个冠军
zezexihaha6 小时前
从“帮写文案”到“管生活”:个人AI工具的边界在哪?
人工智能
算家云7 小时前
nano banana官方最强Prompt模板来了!六大场景模板详解
人工智能·谷歌·ai大模型·算家云·ai生图·租算力,到算家云·nano banana 提示词
暴躁的大熊7 小时前
AI助力决策:告别生活与工作中的纠结,明析抉择引领明智选择
人工智能
Gyoku Mint7 小时前
提示词工程(Prompt Engineering)的崛起——为什么“会写Prompt”成了新技能?
人工智能·pytorch·深度学习·神经网络·语言模型·自然语言处理·nlp