Class ExportSession

export/export.py:110–696

Manages the export process for ExecuTorch models. This class handles the export process through a pipeline of stages: 1. (Optional) Quantize - Apply post-training quantization to the model 2. Export - Export PyTorch model to ExportedProgram 3. EdgeTransformAndLower - Transform

    "This API and all of its related functionality such as ExportSession and ExportRecipe are experimental."
)
class ExportSession:
    """
    Manages the export process for ExecuTorch models.

    This class handles the export process through a pipeline of stages:
    1. (Optional) Quantize - Apply post-training quantization to the model
    2. Export - Export PyTorch model to ExportedProgram
    3. EdgeTransformAndLower - Transform and lower to EdgeProgramManager
    4. Executorch - Convert to ExecutorchProgramManager for final execution
    """

    def __init__(
        self,
        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: Optional[bool] = False,
    ) -> None:
        """
        Initialize the ExportSession with model, inputs, and recipe.

        Args:
            model: The PyTorch model(s) to export. Can be:
                  - nn.Module or Dict[str, nn.Module]: Eager PyTorch model(s)
                  - GraphModule or Dict[str, GraphModule]: Quantized model(s)
                  - ExportedProgram or Dict[str, ExportedProgram]: Already exported model(s)
                  - str: Path to load an ExportedProgram from disk
            example_inputs: Example inputs for the model(s), either a list of input tuples
                          or a dictionary mapping method names to lists of input tuples.
                          First sample (index 0) is used for torch.export.export() to export the model.
                          All samples are used as calibration dataset in PT2E Quantize stage,
                          Optional when model is ExportedProgram (not needed).
            export_recipe: Contains the configuration for the export process
            name: Optional name for the export
            dynamic_shapes: Optional dynamic shape specifications
            constant_methods: Optional dictionary of constant methods
            artifact_dir: Optional directory to store artifacts
            generate_etrecord: Optional flag to generate an etrecord
        """
        # Load model from file if string path provided
        if isinstance(model, str):