Model Performance

Strategies for improving inference performance through data shape optimizations.

Data optimizations are strategies that yield inference results with faster speed and/or less resources by optimizing the format of the data or how it is transmitted. It is important to note that inputs and outputs can and should use the optimization strategy appropriate for each.

The following shows different input schema options and the performance gains for each optimization. These image processing samples have large arrays in the shape (2052, 2456). The following is a sample based on converting images and performing inferences with models in the Wallaroo BYOP Framework with example conversions.

Data Type	Performance Benchmark	Code Snippet
Unmodified Pandas DataFrame	100 seconds	`input_data = pd.DataFrame({"image": [image.tolist()]})`
Unmodified Apache Arrow	10 Seconds	`input_data = pa.Table.from_pandas(df, schema=input_schema)`
Apache Arrow with Flattened Arrays	65 ms	`input_data = pa.Table.from_pydict({` `'image': pa.FixedSizeListArray.from_arrays(image.ravel(), len(image.ravel())),` `'dim0': [image.shape[0]],` `'dim1': [image.shape[1]],` `})`
Apache Arrow with Fixed Shape Tensor Arrays	18 ms	`input_data = pa.Table.from_pydict({` `'image': pa.FixedShapeTensorArray.from_numpy_ndarray(image),` `})`

The method of data optimization depends on what use case the model is designed for and the data type of the inputs and outputs. The following options are provided as guidelines for how to structure the data based on use cases. The following table provides recommendations based on the model’s use case.

Data Types	Use Case	Recommended Data Optimization
Fixed Array Shape	Static Image Tasks	Apache Arrow with Fixed Shape Tensor Arrays
Variable Array Shape	Video and Audio Analysis	Apache Arrow with Data Flattening.
Strings	Text Analysis	Unmodified Pandas DataFrame. Note in this instances, no modifications to the ML Model or BYOP implementation are required.
Large File Sizes with Data Types Not Listed Above	Any	Unmodified Apache Arrow. Note in this instances, no modifications to the ML Model or BYOP implementation are required.
Inputs to MLFlow Models	MLFlow Models	Unmodified Pandas DataFrame. MLFlow models only accept pandas DataFrames.

For use cases where the data shape is fixed such as static images of the same size and resolution, the optimization recommendation is to use Apache Arrow Fixed Shape Tensor Arrays or pyarrow.FixedSizeListArray. The following example demonstrates a Wallaroo BYOP Framework code to support this data scheme.

Note that both the input and output schemas use the fix shape tensor data type, but they are not required to be the same data type or use same optimization.

The model input and output schema are as follows.

# input schema

import pyarrow as pa

input_schema = pa.schema([
    pa.field('image', pa.fixed_shape_tensor(pa.uint16(), [2052, 2456])),
])

# output schema

output_schema = pa.schema([
    pa.field('image', pa.fixed_shape_tensor(pa.uint16(), [2052, 2456])),
    pa.field('virtual_stain', pa.fixed_shape_tensor(pa.uint8(), [2052, 2456, 3])),
    pa.field('mask_overlay', pa.fixed_shape_tensor(pa.uint8(), [2052, 2456, 3])),
    pa.field('mask', pa.fixed_shape_tensor(pa.uint16(), [2052, 2456])),
])

The sample BYOP code to accept this data is:

import logging

import numpy as np
from mac.types import InferenceData


def process_data(input_data: InferenceData) -> InferenceData:
    """Dummy processing step."""
    logging.info(f"Got keys: {input_data.keys()}")
    image_3d = np.zeros((1, 2052, 2456, 3))
    image_2d = np.zeros((1, 2052, 2456))

    return {
        # set return data here based on the output
    }

Use cases where the data shape is variable, the recommended data optimization is Apache Arrow with Data Flattening. In these instances, multi-dimensional arrays are flattened into one array. To reconstruct the original shape, additional parameters are provided to inform Wallaroo of the original shape so it can be reconstituted.

In the following example, the variables dim0 and dim1 are used to indicate the shape of the data to convert to from the flattened array.

# input schema
input_schema = pa.schema([
    pa.field('image', pa.list_(pa.uint16())),
    pa.field('dim0', pa.int64()),
    pa.field('dim1', pa.int64()),
])

The following example is the Wallaroo BYOP Framework code that supports this data scheme.

import logging

import numpy as np
from mac.types import InferenceData


def process_data(input_data: InferenceData) -> InferenceData:
    """Dummy processing step."""
    logging.info(f"Got keys: {input_data.keys()}")
    dim0 = input_data["dim0"][0]
    dim1 = input_data["dim1"][0]
    shape = (dim0, dim1)
    bf_image = input_data["image"].reshape(shape)
    logging.info(f"Got bf_image: {bf_image.shape if bf_image is not None else None})")
    image_3d = np.zeros((1, 2052, 2456, 3))
    image_2d = np.zeros((1, 2052, 2456))

    return {
        # set return data here based on the output
    }

For use cases where the data sizes are large and the data shape or model type does not match the previous examples, converting to Apache Arrow provides an immense improvement in data transmission and inference speed.

As an example, the following image converted to pandas DataFrame is 11.9 MB in size, while the Apache Arrow Table is only 7.4 MB. As seen in the table above, this can improve inference time dramatically.

The following example shows the data input and output schema for a model that accepts both pandas DataFrames and Apache Arrow Tables.

# input schema
input_schema = pa.schema([
    pa.field('image', pa.list_(pa.list_(pa.uint16()))),
])

# output schema
output_schema = pa.schema([
    pa.field('image', pa.list_(pa.list_(pa.uint16()))),
    pa.field('virtual_stain', pa.list_(pa.list_(pa.list_(pa.uint8())))),
    pa.field('mask_overlay', pa.list_(pa.list_(pa.list_(pa.uint8())))),
    pa.field('mask', pa.list_(pa.list_(pa.uint16()))),
])

The following example is the Wallaroo BYOP Framework code that supports this data scheme.

import logging

import numpy as np
from mac.types import InferenceData


def process_data(input_data: InferenceData) -> InferenceData:
    """Dummy processing step."""
    logging.info(f"Got keys: {input_data.keys()}")

    return {
        "image": input_data['image']
    }

For String based inputs, for example, Large Language Models (LLMs) or MLFLow models, the recommended data optimization is pandas DataFrames.