【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

Logo

开放原子开发者工作坊旨在鼓励更多人参与开源活动,与志同道合的开发者们相互交流开发经验、分享开发心得、获取前沿技术趋势。工作坊有多种形式的开发者活动,如meetup、训练营等,主打技术交流,干货满满,真诚地邀请各位开发者共同参与!

更多推荐