PAI-Blade Python API signature, input parameters, and return value - Platform For AI

PAI-Blade exposes an optimize() function that integrates model optimization directly into your Python pipeline. Pass in a TensorFlow or PyTorch model, choose an optimization level and target device, and get back an optimized model ready for inference.

optimize

def optimize(
    model: Any,
    optimization_level: str,
    device_type: str,
    config: Optional[Config] = None,
    inputs: Optional[List[str]] = None,
    outputs: Optional[List[str]] = None,
    input_shapes: Optional[List[List[str]]] = None,
    input_ranges: Optional[List[List[str]]] = None,
    test_data: List[Dict[str, np.ndarray]] = [],
    calib_data: List[Dict[str, np.ndarray]] = [],
    custom_ops: List[str] = [],
    verbose: bool = False,
) -> Tuple[Any, OptimizeSpec, OptimizeReport]:
    pass

Quick start

The following example shows the minimal call to optimize a TensorFlow SavedModel for GPU inference using lossless optimization:

import blade

optimized_model, opt_spec, opt_report = blade.optimize(
    model="/path/to/saved_model",   # path to a TensorFlow SavedModel directory
    optimization_level="o1",        # lossless optimization: graph rewriting + compilation
    device_type="gpu",
)

To use quantization (o2) with calibration data:

import blade
import numpy as np

# Prepare calibration data as a list of feed_dict arguments
calib = [
    {"input": np.random.rand(1, 224, 224, 3).astype(np.float32)},
    {"input": np.random.rand(1, 224, 224, 3).astype(np.float32)},
]

optimized_model, opt_spec, opt_report = blade.optimize(
    model="/path/to/saved_model",
    optimization_level="o2",   # quantization — requires calib_data
    device_type="gpu",
    calib_data=calib,
)

Parameters

Required parameters

Parameter	Type	Description
`model`	Multiple types	The model to optimize. See Supported model formats.
`optimization_level`	string	The optimization strategy. Valid values (case-insensitive): `o1`, `o2`. See optimization_level values.
`device_type`	string	The target device. Valid values (case-insensitive): `gpu`, `cpu`, `edge` (TensorFlow only).

Optional parameters

Parameter	Type	Default	Description
`inputs`	list[string]	`None`	Names of input nodes. If not set, PAI-Blade infers them automatically.
`outputs`	list[string]	`None`	Names of output nodes. If not set, PAI-Blade infers them automatically.
`input_shapes`	list[list[string]]	`None`	Possible input tensor shapes. Improves optimization for specific shape patterns. See input_shapes format.
`input_ranges`	list[list[string]]	`None`	Value ranges of input tensors. See input_shapes format.
`test_data`	Multiple types	`[]`	Test data used to calibrate the run speed of the model. Data type matches `calib_data` format.
`calib_data`	Multiple types	`[]`	Calibration data for quantization. Required when `optimization_level` is `o2`. Data type depends on the model framework.
`custom_ops`	list[string]	`[]`	Paths to custom operator libraries. Add each library path if the model depends on custom operators.
`verbose`	bool	`False`	Set to `True` to print detailed optimization logs. Useful for diagnosing issues.
`config`	blade.Config	`None`	Advanced optimization settings. See blade.Config.

Supported model formats

TensorFlow models accept:

A GraphDef object
A path to a GraphDef PB file (.pb or .pbtxt)
A path to a SavedModel directory

PyTorch models accept:

A torch.nn.Module object
A path to an exported torch.nn.Module file (.pt)

optimization_level values

Value	Strategy	When to use
`o1`	Lossless optimization — graph rewriting and compilation optimization	Use when accuracy must be preserved exactly. No calibration data needed. Recommended as the starting point.
`o2`	Quantization	Use when you want maximum inference speedup. Requires `calib_data`.

calib_data is only meaningful when optimization_level is o2. It has no effect for o1.

input_shapes and input_ranges format

Both parameters use the same nested list structure: each inner list holds one value per input tensor, and the outer list holds multiple shape groups.

input_shapes — each element is a string in "<dim1>*<dim2>*..." format:

# Single model with two inputs, one shape group
input_shapes = [["1*512", "3*256"]]

# Single model with two inputs, three shape groups
input_shapes = [
    ["1*512", "3*256"],
    ["5*512", "9*256"],
    ["10*512", "27*256"],
]

input_ranges — each element is a string specifying a range using brackets, real numbers, or characters:

# Single model with two inputs, one range group
input_ranges = [["[0.1,0.4]", "[a,f]"]]

# Single model with two inputs, three range groups
input_ranges = [
    ["[0.1,0.4]", "[a,f]"],
    ["[1.1,1.4]", "[h,l]"],
    ["[2.1,2.4]", "[n,z]"],
]

calib_data format

The data type of calib_data (and test_data) depends on the model framework:

Framework	Type	Description
TensorFlow	`list[dict[string, np.ndarray]]`	A list of `feed_dict` arguments, one per calibration sample
PyTorch	`list[tuple[torch.tensor, ...]]`	A list of input tensor tuples, one per calibration sample

Return value

optimize() returns a tuple of three elements: Tuple[Any, OptimizeSpec, OptimizeReport].

Position	Element	Type	Description
1	Optimized model	Multiple types	Same type as the input model. For example, if you pass a TensorFlow SavedModel directory, a `GraphDef` object is returned.
2	External dependencies	`OptimizeSpec`	Environment variables and compilation cache required for the optimization to take effect. Apply them using the `WITH` statement in Python. Not required when using the SDK.
3	Optimization report	`OptimizeReport`	Details on optimization results. For more information, see Optimization report.

blade.Config

Use blade.Config to fine-tune advanced optimization behavior. Pass an instance to the config parameter of optimize().

class Config(ABC):
    def __init__(
        self,
        disable_fp16_accuracy_check: bool = False,
        disable_fp16_perf_check: bool = False,
        enable_static_shape_compilation_opt: bool = False,
        enable_dynamic_shape_compilation_opt: bool = True,
        quant_config: Optional[Dict[str, str]] = None,
    ) -> None:
        pass

Example:

import blade

config = blade.Config(
    enable_static_shape_compilation_opt=True,
    quant_config={"weight_adjustment": "true"},
)

optimized_model, opt_spec, opt_report = blade.optimize(
    model="/path/to/saved_model",
    optimization_level="o2",
    device_type="gpu",
    calib_data=calib,
    config=config,
)

blade.Config parameters

Parameter	Type	Default	Description
`disable_fp16_accuracy_check`	bool	`False`	Specifies whether to enable accuracy verification in FP16 optimization. `False` (default): disables accuracy verification. `True`: enables accuracy verification.
`disable_fp16_perf_check`	bool	`False`	Specifies whether to enable performance verification in FP16 optimization. `False` (default): disables performance verification. `True`: enables performance verification.
`enable_static_shape_compilation_opt`	bool	`False`	Enables static shape compilation optimization. Set to `True` if all input shapes are fixed at inference time.
`enable_dynamic_shape_compilation_opt`	bool	`True`	Enables dynamic shape compilation optimization. Enabled by default to handle variable-length inputs.
`quant_config`	dict[string, string]	`None`	Quantization settings. Only the `weight_adjustment` key is supported. Set to `"true"` to reduce precision loss by adjusting model parameters during quantization, or `"false"` to disable it.