Deletes all unused submodules from ``self``. A Module is considered "used" if any one of the following is true: 1. It has children that are used 2. Its forward is called directly via a ``call_module`` node 3. It has a non-Module attribute that is use
(self)
| 637 | |
| 638 | @compatibility(is_backward_compatible=True) |
| 639 | def delete_all_unused_submodules(self) -> None: |
| 640 | """ |
| 641 | Deletes all unused submodules from ``self``. |
| 642 | |
| 643 | A Module is considered "used" if any one of the following is |
| 644 | true: |
| 645 | 1. It has children that are used |
| 646 | 2. Its forward is called directly via a ``call_module`` node |
| 647 | 3. It has a non-Module attribute that is used from a |
| 648 | ``get_attr`` node |
| 649 | |
| 650 | This method can be called to clean up an ``nn.Module`` without |
| 651 | manually calling ``delete_submodule`` on each unused submodule. |
| 652 | """ |
| 653 | used: List[str] = [] |
| 654 | |
| 655 | for node in self.graph.nodes: |
| 656 | |
| 657 | if node.op == "call_module" or node.op == "get_attr": |
| 658 | |
| 659 | # A list of strings representing the different parts |
| 660 | # of the path. For example, `foo.bar.baz` gives us |
| 661 | # ["foo", "bar", "baz"] |
| 662 | fullpath = node.target.split(".") |
| 663 | |
| 664 | # If we're looking at multiple parts of a path, join |
| 665 | # join them with a dot. Otherwise, return that single |
| 666 | # element without doing anything to it. |
| 667 | def join_fn(x: str, y: str) -> str: |
| 668 | return ".".join([x, y] if y else [x]) |
| 669 | |
| 670 | # Progressively collect all the names of intermediate |
| 671 | # modules. For example, if we have the target |
| 672 | # `foo.bar.baz`, we'll add `foo`, `foo.bar`, and |
| 673 | # `foo.bar.baz` to the list. |
| 674 | for path in itertools.accumulate(fullpath, join_fn): |
| 675 | used.append(path) |
| 676 | |
| 677 | # For a `call_module` node, also register all recursive submodules |
| 678 | # as used |
| 679 | if node.op == "call_module": |
| 680 | try: |
| 681 | submod = self.get_submodule(node.target) |
| 682 | |
| 683 | for submod_name, _ in submod.named_modules(): |
| 684 | if submod_name != "": |
| 685 | used.append(".".join([node.target, submod_name])) |
| 686 | except AttributeError: |
| 687 | # Node referenced nonexistent submodule, don't need to |
| 688 | # worry about GCing anything |
| 689 | pass |
| 690 | |
| 691 | to_delete = [name for name, _ in self.named_modules() if name not in used] |
| 692 | |
| 693 | for name in to_delete: |
| 694 | self.delete_submodule(name) |
| 695 | |
| 696 | @property |