get the named module in `mod` by the attribute name, for example ``look_up_named_module(net, "features.3.1.attn")`` Args: name: a string representing the module attribute. mod: a pytorch module to be searched (in ``mod.named_modules()``). print_all_options: whet
(name: str, mod, print_all_options=False)
| 114 | |
| 115 | |
| 116 | def look_up_named_module(name: str, mod, print_all_options=False): |
| 117 | """ |
| 118 | get the named module in `mod` by the attribute name, |
| 119 | for example ``look_up_named_module(net, "features.3.1.attn")`` |
| 120 | |
| 121 | Args: |
| 122 | name: a string representing the module attribute. |
| 123 | mod: a pytorch module to be searched (in ``mod.named_modules()``). |
| 124 | print_all_options: whether to print all named modules when `name` is not found in `mod`. Defaults to False. |
| 125 | |
| 126 | Returns: |
| 127 | the corresponding pytorch module's subcomponent such as ``net.features[3][1].attn`` |
| 128 | """ |
| 129 | name_str = look_up_option( |
| 130 | name, {n[0] for n in mod.named_modules()}, default=None, print_all_options=print_all_options |
| 131 | ) |
| 132 | if name_str is None: |
| 133 | return None |
| 134 | if name_str == "": |
| 135 | return mod |
| 136 | for n in name_str.split("."): |
| 137 | if n.isdigit(): |
| 138 | mod = mod[int(n)] |
| 139 | else: |
| 140 | n = look_up_option(n, {item[0] for item in mod.named_modules()}, default=None, print_all_options=False) |
| 141 | if n is None: |
| 142 | return None |
| 143 | mod = getattr(mod, n) |
| 144 | return mod |
| 145 | |
| 146 | |
| 147 | def set_named_module(mod, name: str, new_layer): |
searching dependent graphs…