onnx_diagnostic.torch_export_patches.onnx_export_serialization¶
- onnx_diagnostic.torch_export_patches.onnx_export_serialization.register_cache_serialization(patch_transformers: bool = False, patch_diffusers: bool = True, verbose: int = 0) Dict[str, bool] [source][source]¶
Registers many classes with
onnx_diagnostic.torch_export_patches.onnx_export_serialization.register_class_serialization()
. Returns information needed to undo the registration.- Parameters:
patch_transformers – add serialization function for transformers package
patch_diffusers – add serialization function for diffusers package
verbosity – verbosity level
- Returns:
information to unpatch
- onnx_diagnostic.torch_export_patches.onnx_export_serialization.register_class_serialization(cls, f_flatten: Callable, f_unflatten: Callable, f_flatten_with_keys: Callable, f_check: Callable | None = None, verbose: int = 0) bool [source][source]¶
Registers a class. It can be undone with
onnx_diagnostic.torch_export_patches.onnx_export_serialization.unregister_class_serialization()
.- Parameters:
cls – class to register
f_flatten – see
torch.utils._pytree.register_pytree_node
f_unflatten – see
torch.utils._pytree.register_pytree_node
f_flatten_with_keys – see
torch.utils._pytree.register_pytree_node
f_check – called to check the registration was successful
verbose – verbosity
- Returns:
registered or not
- onnx_diagnostic.torch_export_patches.onnx_export_serialization.serialization_functions(patch_transformers: bool = False, patch_diffusers: bool = False, verbose: int = 0) Dict[type, Callable[[int], bool]] [source][source]¶
Returns the list of serialization functions.