Python PyTorch gather_object用法及代碼示例

本文簡要介紹python語言中 torch.distributed.gather_object 的用法。

用法:

torch.distributed.gather_object(obj, object_gather_list=None, dst=0, group=None)

參數：

• obj(任何) -輸入對象。必須是 picklable 的。

• object_gather_list(list[任何]) -輸出列表。在dst 等級上，它的大小應該正確地調整為該集合的組大小，並將包含輸出。在非 dst 等級上必須是 None。 (默認為 None )

• dst(int,可選的) -目的地排名。 (默認為 0)

• group-(ProcessGroup，可選)：要處理的進程組。如果沒有，將使用默認進程組。默認為 None 。

返回：

沒有。在dst 等級上，object_gather_list 將包含集體的輸出。

在單個進程中從整個組中收集 picklable 的對象。與gather()類似，但可以傳入Python對象。請注意，該對象必須是可picklable的才能被收集。

注意

請注意，此 API 與聚集集合略有不同，因為它不提供 async_op 句柄，因此將是一個阻塞調用。

注意

請注意，使用 NCCL 後端時不支持此 API。

警告

gather_object() 隱式使用 pickle 模塊，已知這是不安全的。可以構造惡意的 pickle 數據，該數據將在 unpickling 期間執行任意代碼。僅使用您信任的數據調用此函數。

例子：

>>> # Note: Process group initialization omitted on each rank.
>>> import torch.distributed as dist
>>> # Assumes world_size of 3.
>>> gather_objects = ["foo", 12, {1: 2}] # any picklable object
>>> output = [None for _ in gather_objects]
>>> dist.gather_object(
        gather_objects[dist.get_rank()],
        output if dist.get_rank() == 0 else None,
        dst=0
    )
>>> # On rank 0
>>> output
['foo', 12, {1: 2}]