torch分布式训练笔记

**Torch 分布式训练笔记**

在深度学习领域，随着数据集的不断增长和模型的不断复杂化，单机训练已经无法满足需求。分布式训练成为必然选择之一。在本文中，我们将介绍 Torch 的分布式训练相关知识，并提供一些实践示例。

**1. 分布式训练的基本概念**

分布式训练是指利用多台计算机（或 GPU）同时进行模型训练，提高训练效率。每台计算机负责处理一部分数据和模型参数。

**2. Torch 的分布式训练支持**

Torch 提供了强大的分布式训练支持，可以在单机、多机甚至云上进行分布式训练。我们可以使用 `torch.distributed` 模块来实现分布式训练。

**3. 分布式训练的准备工作**

###3.1 硬件准备* 多台计算机（或 GPU），每台至少有一个 GPU。
* 网络环境，确保各个计算机之间可以通信。

###3.2 软件准备* Torch，版本 >=1.9.0。
* `torch.distributed` 模块。

**4. 分布式训练的基本流程**

1. **数据分割**: 将原始数据集分成多个子集，每个子集对应一台计算机。
2. **模型参数初始化**: 在每台计算机上初始化模型参数。
3. **模型训练**: 每台计算机在本地进行模型训练，使用 `torch.distributed` 模块来同步模型参数和梯度。
4. **模型融合**: 将各个子集的模型参数融合起来，得到最终模型。

**5. 实现分布式训练**

###5.1 分布式数据分割

import torch.distributed as dist# 初始化分布式环境dist.init_process_group(backend='nccl', init_method='env://')

# 获取当前进程号rank = dist.get_rank()

# 获取总进程数world_size = dist.get_world_size()

# 数据分割data_list = []
for i in range(world_size):
 data_list.append([i *10, (i +1) *10])

print(f"Rank {rank} 的数据：{data_list[rank]}")

import torch.nn as nnimport torch.distributed as dist# 初始化分布式环境dist.init_process_group(backend='nccl', init_method='env://')

# 获取当前进程号rank = dist.get_rank()

# 获取总进程数world_size = dist.get_world_size()

# 模型定义class Net(nn.Module):
 def __init__(self):
 super(Net, self).__init__()
 self.fc1 = nn.Linear(5,10)
 self.fc2 = nn.Linear(10,20)

 def forward(self, x):
 x = torch.relu(self.fc1(x))
 x = self.fc2(x)
 return x# 模型初始化net = Net()
dist.barrier() # 等待所有进程完成模型初始化# 模型训练criterion = nn.MSELoss()
optimizer = torch.optim.SGD(net.parameters(), lr=0.01)

for epoch in range(10):
 optimizer.zero_grad()

 inputs = torch.randn(100,5)
 labels = torch.randn(100,20)

 outputs = net(inputs)
 loss = criterion(outputs, labels)

 loss.backward()
 optimizer.step()

 if rank ==0:
 print(f"Epoch {epoch +1}, Loss: {loss.item():.4f}")

import torch.distributed as dist# 获取当前进程号rank = dist.get_rank()

# 获取总进程数world_size = dist.get_world_size()

# 模型融合if rank ==0:
 model_list = []
 for i in range(world_size):
 model_list.append(Net())

 # 将各个子集的模型参数融合起来 for i in range(world_size):
 model_list[i].load_state_dict(torch.load(f"model_{i}.pth"))

 #保存最终模型 torch.save(model_list[0].state_dict(), "final_model.pth")
else:
 # 将本地模型参数保存到文件中 torch.save(Net().state_dict(), f"model_{rank}.pth")