Python PyTorch ShardedEmbeddingCollection.named_modules用法及代碼示例

本文簡要介紹python語言中 torchrec.distributed.embedding.ShardedEmbeddingCollection.named_modules 的用法。

用法:

named_modules(memo: Optional[Set[torch.nn.modules.module.Module]] = None, prefix: str = '', remove_duplicate: bool = True) → Iterator[Tuple[str, torch.nn.modules.module.Module]]

參數：

• memo-用於存儲已添加到結果中的模塊集的備忘錄

• prefix-將添加到模塊名稱的前綴

• remove_duplicate-是否刪除結果中重複的模塊實例

生成(Yield)：

(string, Module) - 名稱和模塊的元組

返回網絡中所有模塊的迭代器，產生模塊的名稱以及模塊本身。

注意

重複的模塊隻返回一次。在以下示例中，l 將僅返回一次。

例子：

>>> l = nn.Linear(2, 2)
>>> net = nn.Sequential(l, l)
>>> for idx, m in enumerate(net.named_modules()):
        print(idx, '->', m)

0 -> ('', Sequential(
  (0): Linear(in_features=2, out_features=2, bias=True)
  (1): Linear(in_features=2, out_features=2, bias=True)
))
1 -> ('0', Linear(in_features=2, out_features=2, bias=True))