inference#

Inference-related utility functions for model export and optimizations.

Used for exporting models into various formats such ONNX, torchscript etc. and apply optimizations such as fusion.

Functions

`export_for_inference`	Export a model for inference.
`export_with_logger`	Helper method for exporting a model for inference.

Classes

ExportFormat

Enum class for the supported export formats.

class composer.utils.inference.ExportFormat(value)[source]#

Bases: composer.utils.string_enum.StringEnum

Enum class for the supported export formats.

torchscript#: Export in “torchscript” format.

onnx#: Export in “onnx” format.

composer.utils.inference.export_for_inference(model, save_format, save_path, save_object_store=None, sample_input=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)
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

composer.utils.inference.export_with_logger(model, save_format, save_path, logger, save_object_store=None, sample_input=None, transforms=None)[source]#

Helper method for exporting a model for inference.

Exports the model to: 1) save_object_store, if one is provided, 2) logger.file_artifact(save_path), if (1) does not apply and the logger has a destination that supports file artifact logging, 3) locally, if (1) and (2) do not apply.

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.
logger (Logger) – If this logger has a destination that supports file artifacting logging, and save_object_store is not provided, this logger is used to export the 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 the logger should be used to export the model or 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)
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)

Returns

None