def __init__(self): super(LargeModel, self).__init__() self.fc1 = nn.Linear(784, 1024) self.fc2 = nn.Linear(1024, 1024) self.fc3 = nn.Linear(1024, 10) def forward(self, x): x = self.fc1(x) x = x.cuda(1) # 将部分计算放到GPU 1上 x = self.fc2(x) x = x.cuda(0) # 将部分计算放到GPU 0上 x = self.fc3(x) return x

### PyTorch多GPU环境下的模型计算分配 在PyTorch中，可以通过`torch.cuda.set_device()`来设置单个GPU用于训练[^1]。然而，在多GPU环境中，通常会使用`torch.nn.DataParallel`模块或者更现代的分布式数据并行方法（如`DistributedDataParallel`）。以下是基于`DataParallel`实现多GPU训练的一个典型示例及其解释。 #### 示例代码 ```python import torch import torch.nn as nn from torch.utils.data import DataLoader, TensorDataset # 定义一个简单的神经网络模型 class SimpleNet(nn.Module): def __init__(self): super(SimpleNet, self).__init__() self.fc = nn.Linear(10, 1) def forward(self, x): return torch.sigmoid(self.fc(x)) # 创建虚拟数据集 input_data = torch.randn(100, 10) labels = torch.randint(0, 2, (100,)).float() dataset = TensorDataset(input_data, labels) dataloader = DataLoader(dataset, batch_size=10) # 初始化模型并将模型移动到GPU model = SimpleNet() if torch.cuda.device_count() > 1: # 检查是否有多个可用GPU model = nn.DataParallel(model) # 使用DataParallel封装模型 device = torch.device("cuda" if torch.cuda.is_available() else "cpu") model.to(device) # 设置损失函数和优化器 criterion = nn.BCELoss() optimizer = torch.optim.SGD(model.parameters(), lr=0.01) # 训练循环 for epoch in range(5): # 进行5轮迭代 for inputs, targets in dataloader: inputs, targets = inputs.to(device), targets.unsqueeze(1).to(device) optimizer.zero_grad() # 清零梯度 outputs = model(inputs) # 前向传播 loss = criterion(outputs, targets) # 计算损失 loss.backward() # 反向传播 optimizer.step() # 更新参数 print(f'Epoch {epoch+1}, Loss: {loss.item()}') ``` --- #### 代码解释 1. **模型定义** `SimpleNet`是一个简单的全连接层神经网络模型，其输入维度为10，输出经过Sigmoid激活函数处理后变为概率值。 2. **多GPU检测与模型封装** 如果系统中有超过一个可用的GPU，则通过`nn.DataParallel(model)`将模型封装成可以在多个GPU上运行的形式。这一步的核心作用是自动分割输入数据，并将其分发给不同的GPU进行独立计算[^2]。 3. **设备选择** 利用`torch.device`判断当前是否存在CUDA支持的GPU设备。如果存在，则将模型迁移到GPU；否则，默认使用CPU作为计算资源。 4. **数据加载与前向传播** 数据被送入模型之前先转换至相应的设备（即GPU或CPU），随后完成一次完整的前向传播过程。注意目标标签也需要调整形状以匹配预测输出的格式。 5. **反向传播与参数更新** 在每次批次的数据处理完成后，清空之前的梯度记录并通过链式法则计算新的梯度值。最后利用随机梯度下降算法对权重参数实施更新操作。 6. **性能监控** 打印每一轮次结束后的平均损失情况以便观察收敛趋势。 --- ### 注意事项 - 虽然`DataParallel`能够简化多GPU编程流程，但它可能带来额外开销以及潜在同步问题。对于大规模深度学习项目推荐采用更加高效的解决方案——`DistributedDataParallel`。 - GPU上的某些操作具有异步特性，这意味着即使调用了某个依赖于硬件加速的功能也不意味着它立即被执行完毕。