Function export_module_to_program

test/models/export_program.py:317–354

Exports the module and returns the serialized program data.

(
    module_class: Type[nn.Module],
    external_constants: bool = False,
)

def export_module_to_program(
    module_class: Type[nn.Module],
    external_constants: bool = False,
) -> ExecutorchProgramManager:
    """Exports the module and returns the serialized program data."""
    torch.manual_seed(0)
    # Look for an optional @staticmethod that defines custom trace params.
    export_kwargs: Dict[str, Any] = {}
    if hasattr(module_class, "get_export_kwargs"):
        # pyre-ignore[16]: pyre doesn't know about get_export_kwargs.
        export_kwargs = module_class.get_export_kwargs()
    export_joint = False
    export_state_names = False
    share_mutable_buffers = False
    if hasattr(module_class, "export_joint"):
        # pyre-ignore[16]: pyre just cant figure it out
        export_joint = module_class.export_joint()
    if hasattr(module_class, "export_state_names"):
        # pyre-ignore[16]: pyre just cant figure it out
        export_state_names = module_class.export_state_names()
    if hasattr(module_class, "get_method_names_to_export"):
        # pyre-ignore[16]: pyre just cant figure it out
        methods = module_class.get_method_names_to_export()
    else:
        methods = ["forward"]
    if hasattr(module_class, "share_mutable_buffers"):
        # pyre-ignore[16]: pyre just cant figure it out
        share_mutable_buffers = module_class.share_mutable_buffers()
    module = ExportedModule.export(
        module_class,
        methods,
        export_joint_graph=export_joint,
        external_constants=external_constants,
        export_state_names=export_state_names,
        share_mutable_buffers=share_mutable_buffers,
        **export_kwargs,
    )
    return module.executorch_program


def main() -> None: