给我提供一份pytorch 分布式代码

以下是一个简单的 PyTorch 分布式代码示例，使用了 PyTorch 内置的分布式工具箱（`torch.distributed`）：

import torch
import torch.distributed as dist
from torch.nn.parallel import DistributedDataParallel as DDP

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

# 定义网络
model = torch.nn.Sequential(
    torch.nn.Linear(10, 100),
    torch.nn.ReLU(),
    torch.nn.Linear(100, 1)
)

# 将模型分布式并行化
model = DDP(model)

# 定义损失函数和优化器
criterion = torch.nn.MSELoss()
optimizer = torch.optim.SGD(model.parameters(), lr=0.001)

# 加载数据
train_dataset = torch.utils.data.TensorDataset(torch.randn(100, 10), torch.randn(100, 1))
train_loader = torch.utils.data.DataLoader(train_dataset, batch_size=16)

# 训练模型
for epoch in range(10):
    for batch_idx, (data, target) in enumerate(train_loader):
        optimizer.zero_grad()
        output = model(data)
        loss = criterion(output, target)
        loss.backward()
        optimizer.step()

# 清理分布式环境
dist.destroy_process_group()

在此示例中，我们首先通过 `dist.init_process_group` 初始化了分布式环境，然后定义了一个简单的全连接神经网络模型。我们使用 `torch.nn.parallel.DistributedDataParallel` 将模型分布式并行化，这使得模型可以在多个 GPU 上运行，并实现了自动梯度求解和参数同步。接下来，我们定义了损失函数和优化器。在加载数据后，我们使用 PyTorch 的标准训练循环对模型进行训练。最后，我们使用 `dist.destroy_process_group()` 清理分布式环境。