load() — pytorch Function Reference

Architecture documentation for the load() function in common.py from the pytorch codebase.

Function python DynamoBenchmarks Microbenchmarking calls 4 called by 1

Entity Profile

DynamoBenchmarks→ Microbenchmarking→ load() — pytorch Function Reference

Dependency Diagram

graph TD
  d4bae6c8_efa9_4b5e_53ab_4e403ff2693e["load()"]
  3df8de63_d71a_714c_93a1_1a0e24c0f362["load()"]
  3df8de63_d71a_714c_93a1_1a0e24c0f362 -->|calls| d4bae6c8_efa9_4b5e_53ab_4e403ff2693e
  1eff8423_1f23_d138_6815_07d8dc29a749["_normalize_bench_inputs()"]
  d4bae6c8_efa9_4b5e_53ab_4e403ff2693e -->|calls| 1eff8423_1f23_d138_6815_07d8dc29a749
  8964c48d_7b16_2455_878e_dbe471c4a037["_register_dataclass_output_as_pytree()"]
  d4bae6c8_efa9_4b5e_53ab_4e403ff2693e -->|calls| 8964c48d_7b16_2455_878e_dbe471c4a037
  f258a26f_4fc8_d6ab_1167_c608fca925fb["export()"]
  d4bae6c8_efa9_4b5e_53ab_4e403ff2693e -->|calls| f258a26f_4fc8_d6ab_1167_c608fca925fb
  3df8de63_d71a_714c_93a1_1a0e24c0f362["load()"]
  d4bae6c8_efa9_4b5e_53ab_4e403ff2693e -->|calls| 3df8de63_d71a_714c_93a1_1a0e24c0f362
  style d4bae6c8_efa9_4b5e_53ab_4e403ff2693e fill:#6366f1,stroke:#818cf8,color:#fff

Relationship Graph

Source Code

benchmarks/dynamo/common.py lines 1441–1466

    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]