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.0
        for 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)