export_for_inference#

composer.utils.export_for_inference(model, save_format, save_path, save_object_store=None, sample_input=None, dynamic_axes=None, surgery_algs=None, transforms=None, load_path=None, load_object_store=None, load_strict=False)[source]#

Export a model for inference.

Parameters

  • model (nn.Module) โ€“ An instance of nn.Module. Please note that model is not modified inplace. Instead, export-related transformations are applied to a copy of the model.

  • save_format (Union[str, ExportFormat]) โ€“ Format to export to. Either "torchscript" or "onnx".

  • save_path โ€“ (str): The path for storing the exported model. It can be a path to a file on the local disk,

  • URL (a) โ€“ in a cloud bucket. For example, my_run/exported_model.

  • set (or if save_object_store is) โ€“ in a cloud bucket. For example, my_run/exported_model.

  • name (the object) โ€“ in a cloud bucket. For example, my_run/exported_model.

  • save_object_store (ObjectStore, optional) โ€“ If the save_path is in an object name in a cloud bucket (i.e. AWS S3 or Google Cloud Storage), an instance of ObjectStore which will be used to store the exported model. Set this to None if save_path is a local filepath. (default: None)

  • sample_input (Any, optional) โ€“ Example model inputs used for tracing. This is needed for โ€œonnxโ€ export. The sample_input need not match the batch size you intend to use for inference. However, the model should accept the sample_input as is. (default: None)

  • dynamic_axes (Any, optional) โ€“ Dictionary specifying the axes of input/output tensors as dynamic. May be required for exporting models using older versions of PyTorch when types cannot be inferred.

  • surgery_algs (Union[Callable, Sequence[Callable]], optional) โ€“ Algorithms that should be applied to the model before loading a checkpoint. Each should be callable that takes a model and returns modified model. surgery_algs are applied before transforms. (default: None)

  • transforms (Sequence[Transform], optional) โ€“ transformations (usually optimizations) that should be applied to the model. Each Transform should be a callable that takes a model and returns a modified model. transforms are applied after surgery_algs. (default: None)

  • load_path (str) โ€“ The path to an existing checkpoint file. It can be a path to a file on the local disk, a URL, or if load_object_store is set, the object name for a checkpoint in a cloud bucket. For example, run_name/checkpoints/ep0-ba4-rank0. (default: None)

  • load_object_store (ObjectStore, optional) โ€“ If the load_path is in an object name in a cloud bucket (i.e. AWS S3 or Google Cloud Storage), an instance of ObjectStore which will be used to retreive the checkpoint. Otherwise, if the checkpoint is a local filepath, set to None. (default: None)

  • load_strict (bool) โ€“ Whether the keys (i.e., model parameter names) in the model state dict should perfectly match the keys in the model instance. (default: False)

Returns

None