Python module
engine
You can run an inference with our Python API in just a few lines of code:
- Create an
InferenceSession
. - Load a model with
InferenceSession.load()
, which returns aModel
. - Run the model by passing your input to
Model.execute()
, which returns the output.
That’s it! For more detail, see how to run inference with Python.
InferenceSession
class max.engine.InferenceSession(num_threads: int | None = None, device: ~max.driver.driver.Device = <max.driver.driver.CPU object>, **kwargs)
Manages an inference session in which you can load and run models.
You need an instance of this to load a model as a Model
object.
For example:
session = engine.InferenceSession()
model_path = Path('bert-base-uncased')
model = session.load(model_path)
session = engine.InferenceSession()
model_path = Path('bert-base-uncased')
model = session.load(model_path)
-
Parameters:
num_threads (Optional[int]) – Number of threads to use for the inference session. This parameter defaults to the number of physical cores on your machine.
load()
load(model: str | Path | Any, *, custom_extensions: List[str | Path | Any] | str | Path | Any | None = None, custom_ops_path: str | None = None, input_specs: List[TorchInputSpec] | None = None, weights_registry: dict[str, numpy.ndarray] | None = None) → Model
Loads a trained model and compiles it for inference.
Note: PyTorch models must be in TorchScript format.
-
Parameters:
-
model (Union[str, pathlib.Path, Any]) – Path to a model, or a TorchScript model instance. May be a TorchScript model or an ONNX model.
-
custom_extensions (Optional[CustomExtensionsType]) – The extensions to load for the model. : Supports paths to .mojopkg custom ops, .so custom op libraries for PyTorch and .pt torchscript files for torch metadata libraries. Supports
TorchMetadata
andtorch.jit.ScriptModule
objects for torch metadata libraries without serialization.custom_ops_path: str : The path to your custom ops Mojo package. Deprecated, use
custom_extensions
instead. -
input_specs –
The tensor specifications (shape and data type) for each of the model inputs. This is required when loading serialized TorchScript models because they do not include type and shape annotations.
If the model supports an input with dynamic shapes, use
None
as the dimension size inshape
.For example:
session = engine.InferenceSession()
model = session.load(
"clip-vit.torchscript",
input_specs = [
engine.TorchInputSpec(
shape=[1, 16], dtype=DType.int32
),
engine.TorchInputSpec(
shape=[1, 3, 224, 224], dtype=DType.float32
),
engine.TorchInputSpec(
shape=[1, 16], dtype=DType.int32
),
],
)session = engine.InferenceSession()
model = session.load(
"clip-vit.torchscript",
input_specs = [
engine.TorchInputSpec(
shape=[1, 16], dtype=DType.int32
),
engine.TorchInputSpec(
shape=[1, 3, 224, 224], dtype=DType.float32
),
engine.TorchInputSpec(
shape=[1, 16], dtype=DType.int32
),
],
) -
weights_registry – A mapping from names of model weights’ names to their values. The values are currently expected to be NumPy arrays. Currently, only MAX graph models use this argument.
-
-
Returns:
The loaded model, compiled and ready to execute.
-
Return type:
-
Raises:
RuntimeError – If the path provided is invalid.
Model
class max.engine.Model
A loaded model that you can execute.
Do not instantiate this class directly. Instead, create it with
InferenceSession
.
execute()
execute()
execute(**kwargs: Any) → Dict[str, ndarray | dict | list | tuple]
Executes the model with the provided input and returns the outputs.
For example, if the model has one input tensor named “input”:
input_tensor = np.random.rand(1, 224, 224, 3)
model.execute(input=input_tensor)
input_tensor = np.random.rand(1, 224, 224, 3)
model.execute(input=input_tensor)
-
Parameters:
- args – Currently not supported. You must specify inputs using
kwargs
. - kwargs – The input tensors, each specified with the appropriate tensor name
as a keyword and its value as an
ndarray
. You can find the tensor names to use as keywords frominput_metadata
.
- args – Currently not supported. You must specify inputs using
-
Returns:
A dictionary of output values, each as an
ndarray
,Dict
,List
, orTuple
identified by its output name. -
Return type:
Dict
-
Raises:
- RuntimeError – If the given input tensors’ name and shape don’t match what the model expects.
- TypeError – If the given input tensors’ dtype cannot be cast to what the model expects.
input_metadata
property input_metadata*: List[TensorSpec]*
Metadata about the model’s input tensors, as a list of
TensorSpec
objects.
For example, you can print the input tensor names, shapes, and dtypes:
for tensor in model.input_metadata:
print(f'name: {tensor.name}, shape: {tensor.shape}, dtype: {tensor.dtype}')
for tensor in model.input_metadata:
print(f'name: {tensor.name}, shape: {tensor.shape}, dtype: {tensor.dtype}')
output_metadata
property output_metadata*: List[TensorSpec]*
Metadata about the model’s output tensors, as a list of
TensorSpec
objects.
For example, you can print the output tensor names, shapes, and dtypes:
for tensor in model.ouput_metadata:
print(f'name: {tensor.name}, shape: {tensor.shape}, dtype: {tensor.dtype}')
for tensor in model.ouput_metadata:
print(f'name: {tensor.name}, shape: {tensor.shape}, dtype: {tensor.dtype}')
TensorSpec
class max.engine.TensorSpec(shape: List[int | str | None] | None, dtype: DType, name: str)
Defines the properties of a tensor, including its name, shape and data type.
For usage examples, see Model.input_metadata
.
dtype
property dtype*: DType*
A tensor data type.
name
property name*: str*
A tensor name.
shape
The shape of the tensor as a list of integers.
If a dimension size is unknown/dynamic (such as the batch size), its
value is None
.
TorchInputSpec
class max.engine.TorchInputSpec(shape: List[int | str | None] | None, dtype: DType)
Specifies valid input specification for a TorchScript model.
Before you load a TorchScript model, you must create an instance of this class
for each input tensor, and pass them to the input_specs argument of
InferenceSession.load()
.
For example code, see InferenceSession.load()
.
dtype
property dtype*: DType*
A torch input tensor data type.
shape
The shape of the torch input tensor as a list of integers.
If a dimension size is unknown/dynamic (such as the batch size), the
shape should be None
.
Was this page helpful?
Thank you! We'll create more content like this.
Thank you for helping us improve!