Create and configure an ExportSession with the given parameters. This function provides a convenient way to create an ExportSession and optionally run the export process in one step. Args: model: The PyTorch model(s) to export. Can be: - nn.Module or Dict[str
(
model: Union[
nn.Module,
Dict[str, nn.Module],
GraphModule,
Dict[str, GraphModule],
ExportedProgram,
Dict[str, ExportedProgram],
str,
],
example_inputs: Optional[
Union[
List[tuple[torch.Tensor, ...]],
Dict[str, List[tuple[torch.Tensor, ...]]],
]
] = None,
export_recipe: ExportRecipe = None,
name: Optional[str] = None,
dynamic_shapes: Optional[Union[Any, Dict[str, Any]]] = None,
constant_methods: Optional[Union[Dict[str, Callable]]] = None,
artifact_dir: Optional[str] = None,
generate_etrecord: bool = False,
)
| 40 | "This API and all of its related functionality such as ExportSession and ExportRecipe are experimental." |
| 41 | ) |
| 42 | def export( |
| 43 | model: Union[ |
| 44 | nn.Module, |
| 45 | Dict[str, nn.Module], |
| 46 | GraphModule, |
| 47 | Dict[str, GraphModule], |
| 48 | ExportedProgram, |
| 49 | Dict[str, ExportedProgram], |
| 50 | str, |
| 51 | ], |
| 52 | example_inputs: Optional[ |
| 53 | Union[ |
| 54 | List[tuple[torch.Tensor, ...]], |
| 55 | Dict[str, List[tuple[torch.Tensor, ...]]], |
| 56 | ] |
| 57 | ] = None, |
| 58 | export_recipe: ExportRecipe = None, |
| 59 | name: Optional[str] = None, |
| 60 | dynamic_shapes: Optional[Union[Any, Dict[str, Any]]] = None, |
| 61 | constant_methods: Optional[Union[Dict[str, Callable]]] = None, |
| 62 | artifact_dir: Optional[str] = None, |
| 63 | generate_etrecord: bool = False, |
| 64 | ) -> "ExportSession": |
| 65 | """ |
| 66 | Create and configure an ExportSession with the given parameters. |
| 67 | |
| 68 | This function provides a convenient way to create an ExportSession and |
| 69 | optionally run the export process in one step. |
| 70 | |
| 71 | Args: |
| 72 | model: The PyTorch model(s) to export. Can be: |
| 73 | - nn.Module or Dict[str, nn.Module]: Eager PyTorch model(s) |
| 74 | - GraphModule or Dict[str, GraphModule]: Quantized model(s) (e.g., from prepare/convert) |
| 75 | - ExportedProgram or Dict[str, ExportedProgram]: Already exported model(s) |
| 76 | - str: Path to load an ExportedProgram from disk |
| 77 | example_inputs: Example inputs for the model(s), either a list of input tuples |
| 78 | or a dictionary mapping method names to lists of input tuples. |
| 79 | First sample (index 0) is used for torch.export.export() to export the model. |
| 80 | All samples are used as calibration dataset in PT2E Quantize stage. |
| 81 | Optional when model is ExportedProgram (not needed). |
| 82 | export_recipe: Contains the configuration for the export process |
| 83 | name: Optional name for the export |
| 84 | dynamic_shapes: Optional dynamic shape specifications |
| 85 | constant_methods: Optional dictionary of constant methods |
| 86 | artifact_dir: Optional directory to store artifacts |
| 87 | generate_etrecord: Optional flag to generate an etrecord |
| 88 | |
| 89 | Returns: |
| 90 | A configured ExportSession instance with the export process completed if requested |
| 91 | """ |
| 92 | session = ExportSession( |
| 93 | model=model, |
| 94 | example_inputs=example_inputs, |
| 95 | export_recipe=export_recipe, |
| 96 | name=name, |
| 97 | dynamic_shapes=dynamic_shapes, |
| 98 | constant_methods=constant_methods, |
| 99 | artifact_dir=artifact_dir, |