Method test_submodule_manipulation_API

test/test_fx.py:3056–3200

(self)

        self.assertEqual(ref_batchnorm1d(input), traced(input))

    def test_submodule_manipulation_API(self):
        class C(torch.nn.Module):
            def __init__(self):
                super().__init__()
                self.conv = torch.nn.Conv2d(16, 33, 3, stride=2)
                self.param = torch.nn.Parameter(torch.rand(2, 3))

            def forward(self, x):
                return self.conv(torch.cat([self.param, x]))

        class B(torch.nn.Module):
            def __init__(self):
                super().__init__()
                self.linear = torch.nn.Linear(100, 200)
                self.register_buffer("buf", torch.randn(2, 3))
                self.net_c = C()

            def forward(self, x):
                return self.linear(torch.cat([self.buf, self.net_c(x)]))

        class A(torch.nn.Module):
            def __init__(self):
                super().__init__()
                self.net_b = B()
                self.param = torch.nn.Parameter(torch.rand(2, 3))

            def forward(self, x):
                return self.net_b(x) + self.param

        a = symbolic_trace(A())

        a.add_submodule("net_b.net_c.dropout", torch.nn.Dropout(p=0.2))

        conv = [n for n in a.graph.nodes if n.target == "net_b.net_c.conv"][-1]
        with a.graph.inserting_before(conv):
            with warnings.catch_warnings(record=True) as w:
                dropout = a.graph.call_module(module_name="net_b.net_c.dropout",
                                              args=conv.args)
                self.assertEqual(len(w), 0)

        conv.replace_all_uses_with(dropout)
        a.graph.erase_node(conv)
        a.recompile()

        def module_exists(gm: GraphModule, path: str) -> bool:
            return any(path == name for name, _ in gm.named_modules())

        def parameter_exists(gm: GraphModule, path: str) -> bool:
            return (any(path == name for name, _ in gm.named_parameters())
                    and any(path == name for name in gm.state_dict().keys()))

        def buffer_exists(gm: GraphModule, path: str) -> bool:
            return (any(path == name for name, _ in gm.named_buffers())
                    and any(path == name for name in gm.state_dict().keys()))

        # Test that we added the "dropout" submodule
        self.assertTrue(module_exists(a, "net_b.net_c.dropout"))