Scans the module object and returns an iterable for the :class:`~.Tensor` and :class:`~.Module` attributes that agree with the ``predicate``. For multiple calls of this function with same arguments, the order of objects within the returned iterable is guaranteed to be identic
(
self,
*,
recursive: bool = True,
with_key: bool = False,
with_parent: bool = False,
prefix: Optional[str] = None,
predicate: Callable[[Any], bool] = lambda _: True,
seen: Optional[Set[int]] = None
)
| 151 | return outputs |
| 152 | |
| 153 | def _flatten( |
| 154 | self, |
| 155 | *, |
| 156 | recursive: bool = True, |
| 157 | with_key: bool = False, |
| 158 | with_parent: bool = False, |
| 159 | prefix: Optional[str] = None, |
| 160 | predicate: Callable[[Any], bool] = lambda _: True, |
| 161 | seen: Optional[Set[int]] = None |
| 162 | ) -> Union[Iterable[Any], Iterable[Tuple[str, Any]]]: |
| 163 | """Scans the module object and returns an iterable for the :class:`~.Tensor` |
| 164 | and :class:`~.Module` attributes that agree with the ``predicate``. For multiple |
| 165 | calls of this function with same arguments, the order of objects within the |
| 166 | returned iterable is guaranteed to be identical, as long as all the involved |
| 167 | module objects' ``__dict__`` does not change thoughout those calls. |
| 168 | |
| 169 | Args: |
| 170 | recursive: whether to recursively scan all the submodules. |
| 171 | with_key: whether to yield keys along with yielded objects. |
| 172 | with_parent: whether to yield ``self`` along with yielded objects. |
| 173 | prefix: prefix appended to the yielded keys. |
| 174 | predicate: the predication function applied to scanned objects. |
| 175 | seen: a dict that records whether a module has been traversed yet. |
| 176 | """ |
| 177 | if seen is None: |
| 178 | seen = set([id(self)]) |
| 179 | |
| 180 | module_dict = vars(self) |
| 181 | _prefix = "" if prefix is None else prefix + "." |
| 182 | |
| 183 | for key in sorted(module_dict): |
| 184 | for expanded_key, leaf in _expand_structure(key, module_dict[key]): |
| 185 | leaf_id = id(leaf) |
| 186 | if leaf_id in seen: |
| 187 | continue |
| 188 | seen.add(leaf_id) |
| 189 | |
| 190 | if predicate(leaf): |
| 191 | if with_key and with_parent: |
| 192 | yield _prefix + expanded_key, leaf, self |
| 193 | elif with_key: |
| 194 | yield _prefix + expanded_key, leaf |
| 195 | elif with_parent: |
| 196 | yield leaf, self |
| 197 | else: |
| 198 | yield leaf |
| 199 | |
| 200 | if recursive and isinstance(leaf, Module): |
| 201 | yield from leaf._flatten( |
| 202 | recursive=recursive, |
| 203 | with_key=with_key, |
| 204 | with_parent=with_parent, |
| 205 | prefix=_prefix + expanded_key if with_key else None, |
| 206 | predicate=predicate, |
| 207 | seen=seen, |
| 208 | ) |
| 209 | |
| 210 | def parameters(self, recursive: bool = True, **kwargs) -> Iterable[Parameter]: |