コード例 #1
0
    def test_normalize_args(self):
        m = resnet18()

        class FunctionalTracer(torch.fx.Tracer):
            def is_leaf_module(self, m: torch.nn.Module,
                               module_qualified_name: str) -> bool:
                # `leaves` contains the set of standard `nn.Modules` that are not
                # currently symbolically traceable. Ideally this set would be empty
                leaves = set([torch.nn.BatchNorm2d])
                return type(m) in leaves

        traced = torch.fx.GraphModule(m, FunctionalTracer().trace(m))

        input = torch.randn(5, 3, 224, 224)
        ref_outs = traced(input)

        traced = NormalizeArgs(traced).transform()

        test_outs = traced(input)
        self.assertEqual(test_outs, ref_outs)

        modules = dict(traced.named_modules())
        for node in traced.graph.nodes:
            if node.op == 'call_function' and node.target.__module__ == 'torch.nn.functional':
                self.assertEqual(len(node.args), 0)
            if node.op == 'call_module':
                submod_class = modules[node.target].__class__
                nn_class = getattr(torch.nn, submod_class.__name__)
                if submod_class == nn_class:
                    self.assertEqual(len(node.args), 0)
コード例 #2
0
    def test_normalize_modules_exhaustive(self):
        """
        Exhaustively test `NormalizeArgs` on all standard
        torch.nn Module classes
        """
        for test_params in module_tests + new_module_tests:
            if 'constructor' not in test_params:
                constructor = getattr(torch.nn, test_params['module_name'])
            else:
                constructor = test_params['constructor']

            if 'constructor_args' not in test_params:
                args = ()
            else:
                args = test_params['constructor_args']

            mod = constructor(*args)
            # Skip modules that are not standard `torch.nn`
            # instances, including functionals. (functionals
            # are tested in test_normalize_args)
            if mod.__class__.__name__ not in dir(torch.nn):
                continue

            if 'input_fn' not in test_params:
                inputs = torch.randn(test_params['input_size'])
            else:
                inputs = test_params['input_fn']()

            if not isinstance(inputs, (tuple, list)):
                inputs = (inputs, )

            params = ', '.join(f'v{i}' for i in range(len(inputs)))

            # Generate a class to wrap this standard `nn.Module` instance
            test_classname = f'Test{mod.__class__.__name__}'
            test_mod_code = f"""
class {test_classname}(torch.nn.Module):
    def __init__(self, mod):
        super().__init__()
        self.mod = mod

    def forward(self, {params}):
        return self.mod({params})
            """

            gbls = {'torch': torch}
            exec(test_mod_code, gbls)

            test_instance = gbls[test_classname](mod)
            traced = symbolic_trace(test_instance)

            # Now actually test arg normalization!
            traced = NormalizeArgs(traced).transform()

            # These Modules have an RNG in their forward, so testing
            # correctness by comparing outputs is not correct. Skip that
            # check for these
            stochastic_modules = {
                'FractionalMaxPool2d', 'FractionalMaxPool3d', 'RReLU'
            }

            if mod.__class__.__name__ not in stochastic_modules:
                self.assertEqual(traced(*inputs), mod(*inputs))

            # Ensure all args/kwargs are normalized into kwargs
            modules = dict(traced.named_modules())
            for node in traced.graph.nodes:
                if node.op == 'call_module':
                    submod_class = modules[node.target].__class__
                    nn_class = getattr(torch.nn, submod_class.__name__)
                    if submod_class == nn_class:
                        self.assertEqual(len(node.args), 0)