speechbrain.utils.distributed 模块

仅在主进程上运行某些操作的防护

Authors:

阿卜杜勒·赫巴 2020
阿库·柔赫 2020
彼得·普兰廷加 2023
阿德尔·穆门 2024

摘要

类：

MainProcessContext

上下文管理器，确保代码仅在主进程中运行。

函数：

`ddp_barrier`	在分布式数据并行（DDP）模式下同步所有进程。
`ddp_broadcast`	在DDP模式下，此函数将向所有进程广播一个对象。
`ddp_init_group`	如果python命令行中给出了distributed_launch布尔值，此函数将初始化ddp组。
`get_rank`	获取当前进程的排名。
`if_main_process`	返回当前进程是否为主进程。
`is_distributed_initialized`	返回当前系统是否为分布式系统。
`main_process_only`	函数装饰器，确保函数仅在主进程上运行。
`rank_prefixed_message`	在消息前加上进程的排名。
`run_on_main`	运行一个支持DPP（多GPU）的函数。

参考

speechbrain.utils.distributed.rank_prefixed_message(message: str) → str[source]

在消息前加上进程的等级。

Parameters:: 消息 (str) – 要添加前缀的消息。
Returns:: 如果已知，消息前面会带有等级前缀。
Return type:: str

speechbrain.utils.distributed.get_rank() → int | None[source]

获取当前进程的排名。

这段代码取自Pytorch Lightning库： https://github.com/Lightning-AI/pytorch-lightning/blob/bc3c9c536dc88bfa9a46f63fbce22b382a86a9cb/src/lightning/fabric/utilities/rank_zero.py#L39-L48

Returns:: 当前进程的排名，如果无法确定排名则为None。
Return type:: int 或 None

speechbrain.utils.distributed.run_on_main(func, args=None, kwargs=None, post_func=None, post_args=None, post_kwargs=None, run_post_on_main=False)[source]

运行一个支持DPP（多GPU）的函数。

主函数仅在主进程上运行。可以指定一个post_function，在主函数完成后在非主进程上运行。这样，无论主函数产生什么，都可以在其他进程上加载。

Parameters:

func (callable) – 在主进程上运行的函数。
args (list, None) – 传递给函数的位置参数。
kwargs (dict, None) – 传递给函数的关键字参数。
post_func (callable, None) – 在主函数完成后运行的函数。默认情况下仅在非主进程上运行。
post_args (list, None) – 传递给post_func的位置参数。
post_kwargs (dict, None) – 传递给 post_func 的关键字参数。
run_post_on_main (bool) – 是否也在主进程上运行post_func。（默认值：False）

speechbrain.utils.distributed.is_distributed_initialized() → bool[source]: 返回当前系统是否为分布式系统。

speechbrain.utils.distributed.if_main_process() → bool[source]: 返回当前进程是否为主进程。

class speechbrain.utils.distributed.MainProcessContext[source]

基础类：object

上下文管理器确保代码仅在主进程上运行。这对于确保即使main_proc_wrapped_func函数内部引发异常， MAIN_PROC_ONLY全局变量也会减少非常有用。

__enter__()[source]: 输入上下文。增加计数器。

__exit__(exc_type, exc_value, traceback)[source]: 退出上下文。减少计数器。

speechbrain.utils.distributed.main_process_only(function)[source]: 函数装饰器，用于确保函数仅在主进程上运行。这对于诸如保存到文件系统或记录到网络地址等操作非常有用，您只希望它在单个进程上发生。

speechbrain.utils.distributed.ddp_barrier()[source]

在分布式数据并行（DDP）模式下同步所有进程。

此函数会阻塞当前进程的执行，直到分布式组中的所有进程都达到相同的点。它确保在其它所有进程也达到此屏障之前，没有进程会继续前进。如果未使用DDP（即只有一个进程在运行），此函数将不起作用并立即返回。

Return type:: 无

Example

>>> ddp_barrier()
>>> print("hello world")
hello world

speechbrain.utils.distributed.ddp_broadcast(communication_object, src=0)[source]

在DDP模式下，此函数将向所有进程广播一个对象。

Parameters:

communication_object (Any) – 要传递给所有进程的对象。必须是可序列化的。参见 torch.distributed.broadcast_object_list() 的文档
src (int) – 持有要通信对象的排名。

Return type:

在源等级上传递的通信对象。

speechbrain.utils.distributed.ddp_init_group(run_opts)[source]

如果python命令行中给出了distributed_launch bool，此函数将初始化ddp组。

ddp 组将使用 distributed_backend 参数来设置 DDP 通信协议。RANK Unix 变量将用于将子进程注册到 ddp 组。

Parameters:: run_opts (list) – 要解析的参数列表，通常来自 sys.argv[1:]。
Return type:: 无