The result of exporting an nn.Module. Attributes: eager_module: The original nn.Module that was exported. methods: The names of the eager_module methods that were traced. executorch_program: The resulting ExecutorchProgram. exported_program: The resulting Exporte
| 30 | |
| 31 | |
| 32 | class ExportedModule: |
| 33 | """The result of exporting an nn.Module. |
| 34 | |
| 35 | Attributes: |
| 36 | eager_module: The original nn.Module that was exported. |
| 37 | methods: The names of the eager_module methods that were traced. |
| 38 | executorch_program: The resulting ExecutorchProgram. |
| 39 | exported_program: The resulting ExportedProgram. |
| 40 | trace_inputs: The inputs that were used when tracing eager_module. |
| 41 | """ |
| 42 | |
| 43 | def __init__( |
| 44 | self, |
| 45 | eager_module: nn.Module, |
| 46 | methods: Sequence[str], |
| 47 | executorch_program: ExecutorchProgramManager, |
| 48 | exported_program: torch.export.ExportedProgram, |
| 49 | trace_inputs: Sequence, |
| 50 | get_random_inputs_fn: Callable[[], Sequence], |
| 51 | ): |
| 52 | """INTERNAL ONLY: Use ExportedModule.export() instead.""" |
| 53 | self.eager_module: nn.Module = eager_module |
| 54 | self.methods: Sequence[str] = methods |
| 55 | self.executorch_program: ExecutorchProgramManager = executorch_program |
| 56 | self.exported_program: torch.export.ExportedProgram = exported_program |
| 57 | self.trace_inputs: Sequence = trace_inputs |
| 58 | self.__get_random_inputs_fn = get_random_inputs_fn |
| 59 | |
| 60 | def get_random_inputs(self) -> Sequence: |
| 61 | """Returns random inputs appropriate for model inference.""" |
| 62 | return self.__get_random_inputs_fn() |
| 63 | |
| 64 | @staticmethod |
| 65 | def export( |
| 66 | module_class: Type[nn.Module], |
| 67 | methods: Sequence[str] = ("forward",), |
| 68 | ignore_to_out_var_failure: bool = False, |
| 69 | dynamic_memory_planning_mode: DynamicMemoryPlanningMode = DynamicMemoryPlanningMode.UPPER_BOUND, |
| 70 | export_joint_graph: bool = False, |
| 71 | external_constants: bool = False, |
| 72 | export_state_names: bool = False, |
| 73 | share_mutable_buffers: bool = False, |
| 74 | ) -> "ExportedModule": |
| 75 | """ |
| 76 | Creates a new ExportedModule for the specified module class. |
| 77 | |
| 78 | Args: |
| 79 | module_class: The subclass of nn.Module to export. |
| 80 | methods: The names of the module_class methods to trace. |
| 81 | ignore_to_out_var_failure: Whether to ignore the failue when an |
| 82 | functional op does not have an out variant. |
| 83 | dynamic_memory_planning_mode: The dynamic memory planning mode to |
| 84 | use. |
| 85 | """ |
| 86 | |
| 87 | def get_inputs_adapter( |
| 88 | worker_fn: Callable, method: str |
| 89 | ) -> Callable[[], Sequence]: |