当前位置: 首页 > news >正文

Pytorch中cuda相关操作详见和代码示例

在 PyTorch 中,CUDA 是用于加速深度学习计算的重要接口。下面将从 基础概念常用操作代码示例 等方面详解 PyTorch 中的 CUDA 操作。


一、基础概念

1. torch.device

用于指定 tensor 所在的设备(CPU 或 GPU)。

device = torch.device("cuda")  # 默认 cuda:0

2. torch.cuda.is_available()

检查是否有可用的 CUDA 设备。

if torch.cuda.is_available():device = torch.device("cuda")
else:device = torch.device("cpu")

二、常用 CUDA 操作详解

1. 将张量移动到 GPU

x = torch.randn(2, 3)
x = x.to(device)           # 推荐方式
# 或
x = x.cuda()               # 直接转到默认 GPU

2. 创建时指定设备

x = torch.zeros(3, 3, device='cuda')  # 创建于 GPU 上

3. 多 GPU 支持

torch.cuda.device_count()         # 可用 GPU 数量
torch.cuda.current_device()       # 当前设备 ID
torch.cuda.get_device_name(0)     # 设备名称

4. 模型放入 GPU

model = MyModel()
model.to(device)  # 移动整个模型

模型和数据都要在同一设备上,否则报错。


三、完整代码示例

示例 1:基本张量运算 GPU 加速

import torch# 设备选择
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
print(f"Using device: {device}")# 数据创建与计算
a = torch.rand(10000, 10000, device=device)
b = torch.rand(10000, 10000, device=device)c = torch.matmul(a, b)  # 在 GPU 上计算
print(c)

示例 2:GPU 模型训练

import torch
import torch.nn as nn
import torch.optim as optim# 设备设置
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")# 简单模型
model = nn.Linear(10, 1).to(device)# 损失函数和优化器
criterion = nn.MSELoss()
optimizer = optim.SGD(model.parameters(), lr=0.01)# 输入数据
x = torch.randn(32, 10).to(device)
y = torch.randn(32, 1).to(device)# 前向、反向、优化
for epoch in range(100):pred = model(x)loss = criterion(pred, y)optimizer.zero_grad()loss.backward()optimizer.step()if epoch % 10 == 0:print(f"Epoch {epoch}, Loss: {loss.item()}")

四、CUDA 显存管理

清除显存缓存

torch.cuda.empty_cache()

查看显存使用

print(torch.cuda.memory_allocated() / 1024 ** 2)     # 当前已分配显存 MB
print(torch.cuda.memory_reserved() / 1024 ** 2)      # 当前保留显存 MB

五、进阶用法

1. 多卡训练(手动)

model = nn.DataParallel(model)  # 自动分配到多个 GPU
model.to("cuda")

2. 多卡训练(更推荐)

使用 torch.nn.parallel.DistributedDataParallel(适用于大规模训练)。


六、常见错误

错误原因
RuntimeError: Expected all tensors to be on the same device模型和数据不在同一 GPU
CUDA out of memory显存不够,可减少 batch size 或使用 torch.cuda.empty_cache()
AssertionError: Torch not compiled with CUDA enabled安装的是 CPU 版本 PyTorch,需换为 GPU 版本

七、检查安装的 CUDA 和 PyTorch

print(torch.__version__)         # PyTorch 版本
print(torch.version.cuda)        # 编译时 CUDA 版本
print(torch.cuda.get_device_name(0))  # GPU 型号

八、多 GPU 模型训练模板参考示例

以下是一个基于 PyTorch 的多 GPU 模型训练模板,适用于多个 CUDA 设备并使用 DataParallel 实现数据并行训练(适合中小规模训练任务)。


多 GPU 训练模板(使用 torch.nn.DataParallel

import torch
import torch.nn as nn
import torch.optim as optim
from torchvision import datasets, transforms, models
from torch.utils.data import DataLoader# ==== 1. 设备设置 ====
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
num_gpus = torch.cuda.device_count()
print(f"Using {num_gpus} GPUs on device: {device}")# ==== 2. 数据预处理与加载 ====
transform = transforms.Compose([transforms.ToTensor(),transforms.Normalize((0.5,), (0.5,))
])train_dataset = datasets.CIFAR10(root="./data", train=True, download=True, transform=transform)
train_loader = DataLoader(train_dataset, batch_size=128, shuffle=True, num_workers=4, pin_memory=True)# ==== 3. 构建模型 ====
model = models.resnet18(num_classes=10)if num_gpus > 1:model = nn.DataParallel(model)  # 包装为多 GPU 模型model = model.to(device)# ==== 4. 损失函数与优化器 ====
criterion = nn.CrossEntropyLoss()
optimizer = optim.SGD(model.parameters(), lr=0.01, momentum=0.9)# ==== 5. 训练循环 ====
def train(model, loader, criterion, optimizer, epochs):model.train()for epoch in range(epochs):running_loss = 0.0for i, (inputs, labels) in enumerate(loader):inputs, labels = inputs.to(device, non_blocking=True), labels.to(device, non_blocking=True)optimizer.zero_grad()outputs = model(inputs)loss = criterion(outputs, labels)loss.backward()optimizer.step()running_loss += loss.item()if i % 100 == 0:print(f"[Epoch {epoch+1}/{epochs}] Step {i}, Loss: {loss.item():.4f}")print(f"Epoch {epoch+1} finished. Avg Loss: {running_loss / len(loader):.4f}")# ==== 6. 开始训练 ====
train(model, train_loader, criterion, optimizer, epochs=10)

要点说明

模块说明
DataParallel多 GPU 的快速方案,会自动切分 batch,并在多个 GPU 上并行前向和反向传播
pin_memory=True提高数据从主存拷贝到 GPU 的效率
non_blocking=True异步数据传输
model.to(device)模型必须移动到 GPU,且在 DataParallel 之后再 .to()

推荐进阶:使用 DistributedDataParallel(大规模训练更快)

如果你有多个服务器/节点或者非常大的模型(如 GPT、ViT),建议使用:

torch.nn.parallel.DistributedDataParallel(model)
http://www.dtcms.com/a/297605.html

相关文章:

  • 【Chrome】下载chromedriver的地址
  • 线性代数 下
  • Chrome(谷歌)浏览器 数据JSON格式美化
  • JAVA知识点(六):性能调优与线上问题排查
  • vue+iview+i18n国际化
  • Day 3: 机器学习进阶算法与集成学习
  • 《Uniapp-Vue 3-TS 实战开发》自定义环形进度条组件
  • Zookeeper 3.6.3【详细技术讲解】整
  • Uniapp编写微信小程序,绘制动态圆环进度条
  • Welcome to the world of Go language
  • 鸿蒙端云一体化开发之创建和操作数据库
  • 内存 管理
  • 重读《人件》Peopleware -(22)Ⅲ 适当人选 Ⅵ 乐在其中(上)
  • 微服务架构中的资源调度与负载均衡实践
  • 股指期权可以随时平仓吗?
  • OSPF之多区域
  • cha的操作
  • 每日面试题14:CMS与G1垃圾回收器的区别
  • 2025.07.25【宏基因组】|PathoScope 安装与使用指南
  • Flink 自定义类加载器和子优先类加载策略
  • 编程与数学 03-002 计算机网络 04_数据链路层功能
  • 前端-html+CSS基础到高级(一)html基础
  • centos7安装docker命令
  • Kotlin 数据容器 - List(List 概述、创建 List、List 核心特性、List 元素访问、List 遍历)
  • 论文Review Registration TEASER | TRO | MIT出品,点云配准经典BenchMark | 硬核的数学长文
  • 一文读懂 Doris 冷热分离,优化存储与查询性能
  • Java 大视界 -- Java 大数据机器学习模型在金融衍生品市场波动特征挖掘与交易策略创新中的应用(363)
  • 初识决策树-理论部分
  • [python][flask]flask静态资源
  • OSPF 路由协议多区域