Python PyTorch ShardedEmbeddingBag.named_modules用法及代码示例

本文简要介绍python语言中 torchrec.distributed.embeddingbag.ShardedEmbeddingBag.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))