NativeRTCache Class — pytorch Architecture

Source Code

benchmarks/dynamo/common.py lines 1437–1466

class NativeRTCache:
    cache: dict[weakref.ref, Any] = {}

    @classmethod
    def load(cls, model, example_inputs):
        from torch.export.dynamic_shapes import _combine_args, _tree_map_with_path

        key = weakref.ref(model)
        if key not in cls.cache:
            example_args, example_kwargs = _normalize_bench_inputs(example_inputs)
            example_outputs = model(*example_args, **example_kwargs)
            _register_dataclass_output_as_pytree(example_outputs)

            combined_args = _combine_args(model, example_args, example_kwargs)
            dynamic_shapes = _tree_map_with_path(
                _produce_dynamic_shapes_for_export, combined_args
            )

            ep = torch.export.export(
                model, example_args, example_kwargs, dynamic_shapes=dynamic_shapes
            )
            ep = ep.run_decompositions({})
            with tempfile.NamedTemporaryFile(delete=False) as f:
                torch.export.pt2_archive._package.package_pt2(
                    f, exported_programs={"forward": ep}
                )
                filename = f.name
            cls.cache[key] = PyModelRunner(filename, "forward")

        return cls.cache[key]