在Ubuntu下進行PyTorch的分布式訓練，你需要遵循以下步驟：
安裝PyTorch： 確保你已經安裝了PyTorch。你可以從PyTorch官網根據你的CUDA版本選擇合適的安裝命令。
準備環境： 在開始分布式訓練之前，確保所有參與訓練的機器都已經安裝了相同版本的PyTorch，并且網絡連接正常。
設置環境變量：
為了使分布式訓練正常工作，你需要設置一些環境變量，例如MASTER_ADDR（主節點的IP地址）、MASTER_PORT（一個未被使用的端口號）和WORLD_SIZE（參與訓練的總進程數）。
export MASTER_ADDR='主節點IP'
export MASTER_PORT='端口號'
export WORLD_SIZE='進程總數'
編寫分布式訓練腳本：
在你的PyTorch腳本中，你需要使用torch.distributed包來初始化分布式環境。以下是一個簡單的例子：
import torch
import torch.distributed as dist
from torch.nn.parallel import DistributedDataParallel as DDP
def main(rank, world_size):
# 初始化進程組
dist.init_process_group(
backend='nccl', # 'nccl' for GPU, 'gloo' for CPU
init_method=f'tcp://{MASTER_ADDR}:{MASTER_PORT}',
world_size=world_size,
rank=rank
)
# 創建模型并將其移動到GPU
model = ... # 定義你的模型
model.cuda(rank)
# 使用DistributedDataParallel包裝模型
ddp_model = DDP(model, device_ids=[rank])
# 準備數據加載器
dataset = ... # 定義你的數據集
sampler = torch.utils.data.distributed.DistributedSampler(dataset)
dataloader = torch.utils.data.DataLoader(dataset, batch_size=..., sampler=sampler)
# 訓練循環
for epoch in range(...):
sampler.set_epoch(epoch)
for inputs, targets in dataloader:
inputs, targets = inputs.cuda(rank), targets.cuda(rank)
# 前向傳播
outputs = ddp_model(inputs)
loss = ... # 計算損失
# 反向傳播
loss.backward()
# 更新參數
...
if __name__ == "__main__":
import argparse
parser = argparse.ArgumentParser()
parser.add_argument('--rank', type=int)
parser.add_argument('--world_size', type=int)
args = parser.parse_args()
main(args.rank, args.world_size)
啟動分布式訓練：
使用torch.multiprocessing來啟動多個進程。每個進程都會調用你的訓練腳本，并傳入不同的rank參數。
import torch.multiprocessing as mp
def run(rank, world_size):
main(rank, world_size)
if __name__ == "__main__":
world_size = ... # 總進程數
mp.spawn(run, args=(world_size,), nprocs=world_size, join=True)
運行腳本：
在命令行中，你可以使用mpirun或torch.distributed.launch來啟動分布式訓練。例如：
mpirun -np WORLD_SIZE python your_training_script.py --rank 0
或者使用torch.distributed.launch：
python -m torch.distributed.launch --nproc_per_node=WORLD_SIZE your_training_script.py --rank 0
其中WORLD_SIZE是你的總進程數，--rank是每個進程的排名。
請注意，這些步驟假設你已經有了一個可以分布式訓練的模型和數據集。分布式訓練的具體實現細節可能會根據你的模型和數據集有所不同。此外，確保所有節點之間的SSH無密碼登錄已經設置好，以便于進程間的通信。
