Python torch.multiprocessing.spawn()

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

• 1. 功能与参数说明
  • 2. 示例代码

1. 功能及参数释义

功能：
该模块的核心作用在于生成特定数量的子进程，以并行方式运行目标函数。与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