Python PyTorch DistributedModelParallel.named_buffers用法及代碼示例

本文簡要介紹python語言中 torchrec.distributed.model_parallel.DistributedModelParallel.named_buffers 的用法。

用法:

named_buffers(prefix: str = '', recurse: bool = True) → Iterator[Tuple[str, torch.Tensor]]

參數：

• prefix(str) -前綴添加到所有緩衝區名稱。

• recurse(bool) -如果為真，則產生此模塊和所有子模塊的緩衝區。否則，隻產生作為該模塊直接成員的緩衝區。

生成(Yield)：

(string, torch.Tensor) - 包含名稱和緩衝區的元組

返回模塊緩衝區的迭代器，產生緩衝區的名稱以及緩衝區本身。

例子：

>>> for name, buf in self.named_buffers():
>>>    if name in ['running_var']:
>>>        print(buf.size())