(
model: Union[torch.fx.GraphModule, torch.nn.Module],
example_inputs: Tuple[Value, ...],
*,
example_kwarg_inputs: Optional[Dict[str, Any]] = None,
dynamic_shapes: Optional[Union[Dict[str, Any], Tuple[Any]]] = None,
edge_constant_methods: Optional[Dict[str, Any]] = None,
edge_compile_config=_EDGE_COMPILE_CONFIG,
backend_config=None,
strict=True,
)
| 96 | |
| 97 | |
| 98 | def export_to_exec_prog( |
| 99 | model: Union[torch.fx.GraphModule, torch.nn.Module], |
| 100 | example_inputs: Tuple[Value, ...], |
| 101 | *, |
| 102 | example_kwarg_inputs: Optional[Dict[str, Any]] = None, |
| 103 | dynamic_shapes: Optional[Union[Dict[str, Any], Tuple[Any]]] = None, |
| 104 | edge_constant_methods: Optional[Dict[str, Any]] = None, |
| 105 | edge_compile_config=_EDGE_COMPILE_CONFIG, |
| 106 | backend_config=None, |
| 107 | strict=True, |
| 108 | ) -> ExecutorchProgramManager: |
| 109 | m = model.eval() |
| 110 | # pre-autograd export. eventually this will become torch.export |
| 111 | m = export(m, example_inputs, strict=True).module() |
| 112 | |
| 113 | core_aten_ep = _to_core_aten( |
| 114 | m, |
| 115 | example_inputs, |
| 116 | example_kwarg_inputs=example_kwarg_inputs, |
| 117 | dynamic_shapes=dynamic_shapes, |
| 118 | strict=strict, |
| 119 | ) |
| 120 | |
| 121 | edge_m = _core_aten_to_edge( |
| 122 | core_aten_ep, edge_constant_methods, edge_compile_config |
| 123 | ) |
| 124 | |
| 125 | exec_prog = edge_m.to_executorch(backend_config) |
| 126 | return exec_prog |
| 127 | |
| 128 | |
| 129 | def save_pte_program( |
no test coverage detected