【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)
运行这个示例时,你会看到如下输出(进程的顺序可能不同):
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