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)