- 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, onnx_opset_version=None, load_path=None, load_object_store=None, load_strict=False, input_names=None, output_names=None)#
Export a model for inference.
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_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,
set (or if save_object_store is) – in a cloud bucket. For example,
name (the object) – in a cloud bucket. For example,
save_object_store (ObjectStore, optional) – If the
save_pathis in an object name in a cloud bucket (i.e. AWS S3 or Google Cloud Storage), an instance of
ObjectStorewhich will be used to store the exported model. Set this to
save_pathis a local filepath. (default:
sample_input (Any, optional) – Example model inputs used for tracing. This is needed for “onnx” export. The
sample_inputneed not match the batch size you intend to use for inference. However, the model should accept the
sample_inputas is. (default:
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 None.
surgery_algsare applied before
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.
transformsare applied after
onnx_opset_version (int, optional) – The opset version ONNX should use for exporting. Only used if save_format is
"onnx". Defaults to Pytorch’s default torch.onnx.export opset version, which changes by PyTorch version. (default:
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_storeis set, the object name for a checkpoint in a cloud bucket. For example, run_name/checkpoints/ep0-ba4-rank0. (default:
load_object_store (ObjectStore, optional) – If the
load_pathis in an object name in a cloud bucket (i.e. AWS S3 or Google Cloud Storage), an instance of
ObjectStorewhich will be used to retreive the checkpoint. Otherwise, if the checkpoint is a local filepath, set to
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:
input_names (Sequence[str], optional) – names to assign to the input nodes of the graph, in order. If set to
None, the keys from the sample_input will be used. Fallbacks to
output_names (Sequence[str], optional) – names to assign to the output nodes of the graph, in order. It set to
None, it defaults to