Creates a new ExportedModule for the specified module class. Args: module_class: The subclass of nn.Module to export. methods: The names of the module_class methods to trace. ignore_to_out_var_failure: Whether to ignore the failue when an
(
module_class: Type[nn.Module],
methods: Sequence[str] = ("forward",),
ignore_to_out_var_failure: bool = False,
dynamic_memory_planning_mode: DynamicMemoryPlanningMode = DynamicMemoryPlanningMode.UPPER_BOUND,
export_joint_graph: bool = False,
external_constants: bool = False,
export_state_names: bool = False,
share_mutable_buffers: bool = False,
)
| 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]: |
| 90 | """Returns a function that may bind `method` as a parameter of |
| 91 | `worker_fn`, and ensures that `worker_fn` always returns a list or |
| 92 | tuple. |
| 93 | |
| 94 | Args: |
| 95 | worker_fn: The function to wrap. Must take zero or one |
| 96 | arguments. If it takes one argument, that argument must be |
| 97 | called "method" and expect a string. |
| 98 | method: The name of the method to possibly pass to `worker_fn`. |
| 99 | |
| 100 | Returns: |
| 101 | A function that takes zero arguments and returns a Sequence. |
| 102 | """ |
| 103 | # Names of the parameters of worker_fn. |
| 104 | params = inspect.signature(worker_fn).parameters.keys() |
| 105 | if len(params) == 1: |
| 106 | assert "method" in params, f"Expected 'method' param in {params}" |
| 107 | # Bind our `method` parameter to `worker_fn`, which has the |
| 108 | # signature `func(method: str)`. |
| 109 | worker_fn = functools.partial(worker_fn, method) |
| 110 | else: |
| 111 | assert len(params) == 0, f"Unexpected params in {params}" |
| 112 | # worker_fn takes no parameters. |
| 113 | |
| 114 | def return_wrapper(): |
| 115 | inputs = worker_fn() |
| 116 | # Wrap the return value in a tuple if it's not already a tuple |
| 117 | # or list. |
| 118 | if not isinstance(inputs, (tuple, list)): |
| 119 | inputs = (inputs,) |
| 120 | return inputs |
| 121 | |
| 122 | return return_wrapper |