Model Performance

Strategies for

Data optimizations are strategies that provide inference results with improved speed by optimizing the format of the data or how it is transmitted.

The following shows different input schema options and the performance gains for each option. 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.

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

The method of data optimization depends on what use case the model is designed for or the input type. 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.

Use Case	Data Type	Recommended Data Optimization
Face Tracking and related Static Image Tasks	Fixed Image Sizes	Apache Arrow with Fixed Shape Tensor Arrays
Video and Audio Analysis	Variable Sized Images	Apache Arrow with Data Flattening
Text Analysis	String Inputs	Unmodified Pandas DataFrame
MLFlow Models	Any	Unmodified Pandas DataFrame
Large File Sizes Not Listed Above	Any	Unmodified Apache Arrow

Apache Arrow with Fixed Shape

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.

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 {
        "image": input_data['image'].astype(np.uint16),
        "virtual_stain": image_3d.astype(np.uint8),
        "mask_overlay": image_3d.astype(np.uint8),
        "mask": image_2d.astype(np.uint16),
    }

Apache Arrow with Data Flattening

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 the model of the original shape.

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()),
])

# output schema
output_schema = pa.schema([
    pa.field('image', pa.list_(pa.uint16())),
    pa.field('virtual_stain', pa.list_(pa.uint8())),
    pa.field('mask_overlay', pa.list_(pa.uint8())),
    pa.field('mask', 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()}")
    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 {
        "image": image_2d.reshape(*image_2d.shape[:-2], -1).astype(np.uint16),
        "virtual_stain": image_3d.reshape(*image_3d.shape[:-3], -1).astype(np.uint8),
        "mask_overlay": image_3d.reshape(*image_3d.shape[:-3], -1).astype(np.uint8),
        "mask": image_2d.reshape(*image_2d.shape[:-2], -1).astype(np.uint16),
    }

Unmodified Pandas DataFrame

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

Unmodified Arrow Table

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']
    }