Python PyTorch scatter_object_list用法及代碼示例

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

用法:

torch.distributed.scatter_object_list(scatter_object_output_list, scatter_object_input_list, src=0, group=None)

參數：

• scatter_object_output_list(List[任何]) -非空列表，其第一個元素將存儲分散到該等級的對象。

• scatter_object_input_list(List[任何]) -要分散的輸入對象列表。每個對象都必須是 picklable 的。隻有src rank 上的對象會被分散，非src rank 的參數可以是None。

• src(int) -要從中分散的源排名 scatter_object_input_list 。

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

返回：

None 。如果等級是組的一部分，scatter_object_output_list 將把它的第一個元素設置為這個等級的分散對象。

將scatter_object_input_list中的可pickl對象分散到整個組中。與scatter()類似，但可以傳入Python對象。在每個rank上，分散的對象將存儲為scatter_object_output_list的第一個元素。請注意，scatter_object_input_list 中的所有對象都必須是 picklable 的才能分散。

注意

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

警告

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

例子：

>>> # Note: Process group initialization omitted on each rank.
>>> import torch.distributed as dist
>>> if dist.get_rank() == 0:
>>>     # Assumes world_size of 3.
>>>     objects = ["foo", 12, {1: 2}] # any picklable object
>>> else:
>>>     # Can be any list on non-src ranks, elements are not used.
>>>     objects = [None, None, None]
>>> output_list = [None]
>>> dist.scatter_object_list(output_list, objects, src=0)
>>> # Rank i gets objects[i]. For example, on rank 2:
>>> output_list
[{1: 2}]