1. 使用环境变量选择可见 GPU
在启动程序前，可以通过设置环境变量 CUDA_VISIBLE_DEVICES 来指定哪些 GPU 对程序“可见”。这样做的好处是：
- 无需修改代码：只需要在命令行指定就能控制训练时看到哪些 GPU。
- 重新编号：程序内的 GPU 编号会从 0 开始。例如，如果物理机上 4 块 GPU 的编号为 0、1、2、3，若命令行运行
则程序内部可见的两块 GPU 分别编号为 0（物理机 GPU 1）和 1（物理机 GPU 3）。CUDA_VISIBLE_DEVICES=1,3 python train.py
注：下面的代码示例均没有定义DataLoader
2. 单 GPU 训练
代码示例
import torch
import torch.nn as nn
import torch.optim as optim
# 定义一个简单的神经网络
class SimpleNN(nn.Module):
def __init__(self):
super(SimpleNN, self).__init__()
self.fc1 = nn.Linear(784, 256)
self.relu = nn.ReLU()
self.fc2 = nn.Linear(256, 10)
def forward(self, x):
# 假设输入x的shape为 [batch_size, 784]
x = self.fc1(x)
x = self.relu(x)
x = self.fc2(x)
return x
# 判断GPU是否可用，选择设备
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
print("Using device:", device)
# 将模型放到设备上
model = SimpleNN().to(device)
# 定义损失函数和优化器
criterion = nn.CrossEntropyLoss()
optimizer = optim.Adam(model.parameters(), lr=0.001)
# 模拟训练循环（假设有dataloader提供数据）
for epoch in range(5):
for inputs, labels in dataloader: # dataloader需预先定义好
# 将数据移动到GPU
inputs, labels = inputs.to(device), labels.to(device)
optimizer.zero_grad()
outputs = model(inputs)
loss = criterion(outputs, labels)
loss.backward()
optimizer.step()
print(f"Epoch {epoch+1} complete, Loss: {loss.item():.4f}")
- 设备选择：通过
torch.device("cuda:0")指定使用第 0 块 GPU。若使用环境变量CUDA_VISIBLE_DEVICES限制了 GPU，则“cuda:0”对应的是环境中看到的第一块 GPU。 - 模型移动：使用
model.to(device)将模型参数拷贝到 GPU 上。 - 数据传输：在训练时，将每个 batch 的数据也使用
to(device)移到 GPU 上。
3. 多 GPU 训练：DataParallel
代码示例
import torch
import torch.nn as nn
import torch.optim as optim
class SimpleNN(nn.Module):
def __init__(self):
super(SimpleNN, self).__init__()
self.fc1 = nn.Linear(784, 256)
self.relu = nn.ReLU()
self.fc2 = nn.Linear(256, 10)
def forward(self, x):
x = self.fc1(x)
x = self.relu(x)
x = self.fc2(x)
return x
# 自动检测可用的GPU数量
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
print("GPUs available:", torch.cuda.device_count())
model = SimpleNN()
# 如果有多个GPU则采用DataParallel
if torch.cuda.device_count() > 1:
print("Using DataParallel with", torch.cuda.device_count(), "GPUs")
model = nn.DataParallel(model)
model.to(device)
criterion = nn.CrossEntropyLoss()
optimizer = optim.Adam(model.parameters(), lr=0.001)
for epoch in range(5):
for inputs, labels in dataloader: # 假设已经定义dataloader
inputs, labels = inputs.to(device), labels.to(device)
optimizer.zero_grad()
outputs = model(inputs)
loss = criterion(outputs, labels)
loss.backward()
optimizer.step()
print(f"Epoch {epoch+1} complete, Loss: {loss.item():.4f}")
- DataParallel 原理：
DataParallel 会自动将输入 batch 拆分为多个子 batch，并将模型复制到每个 GPU 上，在各 GPU 上并行计算前向和反向传播，最后将所有 GPU 上计算的梯度汇总到主 GPU 上，再进行参数更新。 - 优点：使用简单，只需一行代码：
model = nn.DataParallel(model)。 - 缺点：存在通信瓶颈（主 GPU 可能负载较高）和显存分配不均的问题，适用于卡数不多、模型较小的场景。
4. 多 GPU 训练：DistributedDataParallel（DDP）
DDP 被认为是多 GPU 训练中最优的方法，尤其适合大规模训练。它采用多进程，每个进程对应一块 GPU，通信采用高效的 AllReduce 操作实现梯度同步。
4.1 DDP 代码示例
下面是一个基于 DDP 的示例代码，假设你已经通过启动脚本（如 torchrun 或 torch.distributed.launch）启动了多个进程。每个进程将使用环境变量 LOCAL_RANK 来确定当前进程使用哪块 GPU。
import os
import torch
import torch.nn as nn
import torch.optim as optim
import torch.distributed as dist
from torch.nn.parallel import DistributedDataParallel as DDP
from torch.utils.data import DataLoader, DistributedSampler
from torchvision import datasets, transforms
def setup_ddp():
# 使用环境变量方式初始化进程组（推荐）
dist.init_process_group(backend='nccl', init_method='env://')
local_rank = int(os.environ["LOCAL_RANK"])
torch.cuda.set_device(local_rank)
return local_rank
class SimpleNN(nn.Module):
def __init__(self):
super(SimpleNN, self).__init__()
self.fc1 = nn.Linear(784, 256)
self.relu = nn.ReLU()
self.fc2 = nn.Linear(256, 10)
def forward(self, x):
x = self.fc1(x)
x = self.relu(x)
x = self.fc2(x)
return x
def main():
local_rank = setup_ddp()
device = torch.device("cuda", local_rank)
# 定义模型，并移动到对应GPU
model = SimpleNN().to(device)
# 使用DistributedDataParallel包装模型
model = DDP(model, device_ids=[local_rank])
criterion = nn.CrossEntropyLoss().to(device)
optimizer = optim.Adam(model.parameters(), lr=0.001)
# 数据预处理与加载（这里使用MNIST示例）
transform = transforms.Compose([
transforms.ToTensor(),
transforms.Lambda(lambda x: x.view(-1)) # 将28x28展平成784维
])
dataset = datasets.MNIST(root='./data', train=True, download=True, transform=transform)
# 使用DistributedSampler确保每个进程处理不同数据
sampler = DistributedSampler(dataset)
dataloader = DataLoader(dataset, batch_size=64, sampler=sampler)
num_epochs = 5
for epoch in range(num_epochs):
sampler.set_epoch(epoch) # 每个epoch打乱顺序
for inputs, labels in dataloader:
inputs, labels = inputs.to(device), labels.to(device)
optimizer.zero_grad()
outputs = model(inputs)
loss = criterion(outputs, labels)
loss.backward()
optimizer.step()
if local_rank == 0:
print(f"Epoch {epoch+1} complete, Loss: {loss.item():.4f}")
# 清理进程组
dist.destroy_process_group()
if __name__ == "__main__":
main()
-
初始化进程组：
调用dist.init_process_group(backend='nccl', init_method='env://')后，所有进程通过环境变量（通常由 torchrun 自动设置）进行初始化。通过LOCAL_RANK确定当前进程应使用哪块 GPU，并用torch.cuda.set_device固定设备。 -
DistributedSampler：
使用DistributedSampler使得每个进程在训练时只读取全数据集的一部分，避免重复计算。记得在每个 epoch 调用sampler.set_epoch(epoch)以确保打乱数据顺序。 -
DDP 包装模型：
通过model = DDP(model, device_ids=[local_rank])将模型封装为分布式版本，这样每个进程仅负责自己的计算，同时在反向传播时自动同步各进程的梯度。 -
启动命令：
使用torchrun（或旧版本的torch.distributed.launch）启动脚本。例如，在单机 4 卡的服务器上，可以使用如下命令：torchrun --nproc_per_node=4 train_ddp.py这样会启动 4 个进程，每个进程自动获得
LOCAL_RANK变量（0~3），分别对应 4 块 GPU。
5. 混合精度训练
混合精度训练能减少显存使用并提高训练速度，尤其在大型模型训练中十分有用。下面是在 DDP 模型中使用混合精度的示例：
from torch.cuda.amp import autocast, GradScaler
# 在模型和优化器初始化后
scaler = GradScaler()
for epoch in range(num_epochs):
sampler.set_epoch(epoch)
for inputs, labels in dataloader:
inputs, labels = inputs.to(device), labels.to(device)
optimizer.zero_grad()
with autocast():
outputs = model(inputs)
loss = criterion(outputs, labels)
scaler.scale(loss).backward()
scaler.step(optimizer)
scaler.update()
if local_rank == 0:
print(f"Epoch {epoch+1} complete, Loss: {loss.item():.4f}")
- autocast：在前向传播时自动使用半精度（FP16），而保持关键计算以单精度（FP32），保证数值稳定。
- GradScaler：用来动态缩放梯度，防止由于使用 FP16 导致梯度过小而出现训练问题。
总结
- 环境变量控制 GPU 可见性：在命令行设置
CUDA_VISIBLE_DEVICES可以灵活选择使用哪几块 GPU。 - 单 GPU 训练：通过
torch.device将模型和数据移动到指定 GPU 上进行训练。 - 多 GPU 训练：
- DataParallel：简单易用，但存在主 GPU 负载过高等问题，适合小规模场景。
- DistributedDataParallel (DDP)：更高效、扩展性更好（支持多机多卡），需要额外设置分布式环境（例如使用 torchrun 启动）。
- 混合精度训练：使用
torch.cuda.amp可以减少显存使用并提高训练速度。 - 数据采样：在 DDP 中使用
DistributedSampler确保各进程数据不重叠并且均衡。
附
显式指定要使用的 GPU
import torch
# 检查 CUDA 是否可用
if torch.cuda.is_available():
# 获取 GPU 数量
num_gpus = torch.cuda.device_count()
print(f"可用 GPU 数量: {num_gpus}")
if num_gpus >= 2:
# 指定使用第二个 GPU (索引为 1)
device = torch.device("cuda:1")
print(f"将使用设备: {device}")
# 将张量或模型移动到指定的 GPU
tensor = torch.randn(10).to(device)
model = YourModel().to(device)
# 你的模型训练代码...
else:
print("系统中没有足够的 CUDA 设备可用。")
else:
print("CUDA 不可用，将使用 CPU。")
使用 torch.nn.DataParallel 或 torch.distributed.DataParallel: 如果你想在多个 GPU 上并行训练你的模型，PyTorch 提供了 torch.nn.DataParallel（适用于单机多 GPU）和 torch.distributed.DataParallel（适用于多机多 GPU）。你可以将你的模型包装在这些类中，PyTorch 会自动将数据分发到不同的 GPU 上进行计算。要使用 DataParallel 并指定要使用的 GPU：
import torch
import torch.nn as nn
class YourModel(nn.Module):
# ... 模型定义 ...
pass
if torch.cuda.is_available():
num_gpus = torch.cuda.device_count()
if num_gpus >= 2:
device_ids = [1] # 指定要使用的 GPU 的索引列表
model = YourModel()
model = nn.DataParallel(model, device_ids=device_ids)
model.to(torch.device(f"cuda:{device_ids[0]}")) # 将模型初始参数放到第一个指定的 GPU 上
# 你的训练代码...
else:
# ...
pass
扫一扫
9711
被折叠的 条评论
为什么被折叠?
到【灌水乐园】发言
