单卡训练改为多卡

### PyTorch 和 TensorFlow 的多 GPU 训练方法 #### 使用 PyTorch 实现多 GPU 分布式训练 在 PyTorch 中，可以通过 `torch.nn.DataParallel` 或更推荐的 `torch.nn.parallel.DistributedDataParallel (DDP)` 来实现多 GPU 并行训练。以下是具体说明： 1. **使用 DataParallel** 数据并行的方式允许模型在多个 GPU 上运行，其中每个 GPU 都会接收到相同的一份模型副本，并处理不同的数据子集。最终的结果会被收集起来并通过平均梯度等方式更新整个模型的参数[^4]。 2. **使用 DDP（DistributedDataParallel）** 更高效的分布式训练方式是采用 DDP 模块。这种方式通过创建多个 Python 进程来管理不同 GPU 上的模型实例，并且需要显式初始化进程组以及设置分布式采样器 `DistributedSampler`。此外，还需要将模型封装到 `DistributedDataParallel` 类中[^2]。 下面是一个简单的代码示例展示如何使用 DDP 完成单机多卡训练： ```python import torch import torch.distributed as dist from torch.nn.parallel import DistributedDataParallel as DDP def setup(rank, world_size): dist.init_process_group(backend='nccl', init_method='env://', rank=rank, world_size=world_size) def cleanup(): dist.destroy_process_group() def main(rank, world_size): setup(rank, world_size) model = YourModel().to(rank) ddp_model = DDP(model, device_ids=[rank]) optimizer = torch.optim.SGD(ddp_model.parameters(), lr=0.001) sampler = torch.utils.data.distributed.DistributedSampler(your_dataset) dataloader = torch.utils.data.DataLoader(dataset, batch_size=batch_size, sampler=sampler) for epoch in range(num_epochs): for data, target in dataloader: output = ddp_model(data.to(rank)) loss = criterion(output, target.to(rank)) optimizer.zero_grad() loss.backward() optimizer.step() cleanup() if __name__ == "__main__": world_size = 4 # 假设有四个GPU torch.multiprocessing.spawn(main, args=(world_size,), nprocs=world_size, join=True) ``` #### 使用 TensorFlow 实现多 GPU 训练 对于 TensorFlow，也可以轻松扩展至多 GPU 场景下工作。通常的做法是在策略层面定义分布模式，比如使用 `tf.distribute.MirroredStrategy()` 来完成同步训练[^3]。 下面是一段基于 MirroredStrategy 的简单例子： ```python import tensorflow as tf strategy = tf.distribute.MirroredStrategy(devices=["/gpu:0", "/gpu:1"]) with strategy.scope(): model = create_your_model() # 创建模型函数 optimizer = tf.keras.optimizers.Adam() dataset = load_and_preprocess_data(batch_size) # 加载和预处理数据函数 distributed_dataset = strategy.experimental_distribute_dataset(dataset) @tf.function def train_step(inputs): features, labels = inputs with tf.GradientTape() as tape: predictions = model(features, training=True) per_example_loss = compute_loss(labels, predictions) # 自定义损失计算函数 total_loss = tf.reduce_sum(per_example_loss) * (1. / global_batch_size) gradients = tape.gradient(total_loss, model.trainable_variables) optimizer.apply_gradients(zip(gradients, model.trainable_variables)) for step, distributed_inputs in enumerate(distributed_dataset): strategy.run(train_step, args=(distributed_inputs,)) ``` 以上展示了两种主流深度学习框架下的多 GPU 设置流程及其核心组件的应用场景。 ---