【python】-——并行训练之torch.multiprocessing.spawn()

1. 功能及参数释义

功能：
主要用于创建指定数量的子进程，并行地执行目标函数。它与torch.distributed.launch的区别是，torch.distributed.launch启动的每个进程， 都运行整个Python 脚本。
主要参数有：

• fn: 要在每个进程中执行的函数。
• args: 传递给 fn 的参数，是一个元组。
• nprocs: 要启动的进程数（通常对应于 GPU 的数量）。
• join: 如果设置为 True，主进程将等待所有子进程结束后再退出。

2. 代码示例

import torch.multiprocessing as mp

def worker(rank, world_size):
    print(f"Hello from rank {rank} out of {world_size} ranks")

if __name__ == "__main__":
    world_size = 4  # 启动四个进程
    mp.spawn(worker, args=(world_size,), nprocs=world_size, join=True)

1
2
3
4
5
6
7
8
9

运行这个示例时，你会看到如下输出（进程的顺序可能不同）：

Hello from rank 0 out of 4 ranks
Hello from rank 1 out of 4 ranks
Hello from rank 2 out of 4 ranks
Hello from rank 3 out of 4 ranks

1
2
3
4
5