Model Schema Optimization
Data optimizations are strategies that provide inference results with improved speed by optimizing the format of the data or how it is transmitted.
The Wallaroo AI inference engine high performance data plane features built-in data processing optimizations that contribute to reducing overall inference roundtrip times for AI models deployed via Wallaroo.
The following shows different input schema options and the performance gains across two data processing options in Wallaroo. 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.
| Pandas DataFrame | Apache Arrow Table | Apache Arrow Table with Data Flattening | Apache Arrow Tablewith Fixed Shape Tensor Arrays | |
|---|---|---|---|---|
| Performance Benchmark | 100 seconds | 10 Seconds | 65 ms | 18 ms |
The data processing optimization method depends on the use case the AI model is designed for as well as its input type/shape. The following methods are provided as guidelines for how to structure the data based on use cases. The following table provides recommendations based on the AI model’s use case.
| Data Shape | Example Use Case(s) | Data Example | Recommended Data Optimization |
|---|---|---|---|
| Variable length, 1-d array | Video and Audio Analysis | Encoded video | Apache Arrow Table |
| Fixed length, 1-d arrays | Time Series Analysis | Variable Sized Images | Apache Arrow Table with Data Flattening |
| Fixed size, N-d arrays | Face Tracking and related Static Image Tasks | Fixed Image Sizes | Apache Arrow Table with Fixed Shape Tensor Arrays |
| Variable sized, N-d arrays | Large File Sizes Not Listed Above | Any | Apache Arrow Table |
Apache Arrow Table
For use cases where the data sizes are large, converting to Apache Arrow provides an immense improvement in data transmission and inference speed.
The sample data is created from a 6144 x 4096 image, with each pixel represented in RGB format, resulting in 3 values per pixel. The model outputs the original image and a virtual stain as the output.
The input and output schemas in Apache pyarrow format are used at model upload to specify the accepted data schemas.
# input schema
list_input_schema = pa.schema([
pa.field('image', pa.list_(pa.list_(pa.list_(pa.uint16())))),
])
# output schema
output_schema = pa.schema([
pa.field('image', pa.list_(pa.uint16())),
pa.field('virtual_stain', pa.list_(pa.uint8()))
])
The image data takes the image and converts it to an multi-dimensional numpy array, with each pixel’s RGB values output.
display(image)
array([[[[ 27, 159, 139],
[114, 107, 101],
[197, 184, 55],
...,
[186, 104, 46],
[254, 67, 40],
[ 96, 180, 188]],
[[ 1, 220, 190],
[ 47, 249, 174],
[ 85, 242, 42],
...,
[167, 32, 11],
[141, 140, 131],
[249, 9, 151]],
[[ 38, 178, 10],
[255, 188, 251],
[106, 62, 216],
...,
[242, 12, 226],
[227, 133, 94],
[ 13, 98, 189]],
...,
...
[118, 197, 154],
...,
[152, 113, 163],
[ 69, 135, 175],
[104, 193, 228]]]], dtype=uint16)
The data is set in an Apache Arrow table in the specified input schema. This is used for the inference input.
arrow_table = pa.Table.from_pydict({"image": [image.tolist()]}, list_input_schema)
display(arrow_table)
pyarrow.Table
image: list<item: list<item: list<item: uint16>>>
child 0, item: list<item: list<item: uint16>>
child 0, item: list<item: uint16>
child 0, item: uint16
----
image: [[[[[26,97,49],[125,62,60],...,[234,216,128],[1,212,234]],[[6,21,220],[25,231,210],...,[221,56,69],[153,52,58]],...,[[70,28,220],[254,231,84],...,[251,160,208],[235,0,143]],[[162,142,105],[159,140,0],...,[26,77,18],[176,143,225]]]]]
The following example is the Wallaroo Custom Model 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:
# //do some processing
# return the original shape and the virtual stain
return {
"image": input_data['image'].astype(np.uint16),
"virtual_stain": image_3d.astype(np.uint8)
}
Apache Arrow Table with Flattened Arrays
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 numpy arrays that are flattened into a single numpy array. To reconstruct the original shape, additional parameters are provided to inform the model of the original shape.
The sample data is created from a 6144 x 4096 image, with each pixel represented in RGB format, resulting in 3 values per pixel. The model outputs the original image and a virtual stain as the output.
The input and output schemas in Apache pyarrow format are used at model upload to specify the accepted data schemas. 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.
For this example, we will return both the image in its submitted flattened state and the shape, with the virtual stain generated from these values.
# 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('dim0', pa.int64()),
pa.field('dim1', pa.int64()),
pa.field('virtual_stain', pa.list_(pa.uint8()))
])
The image data takes the image and converts it to an multi-dimensional numpy array, with each pixel’s RGB values output.
display(image)
array([[[[ 27, 159, 139],
[114, 107, 101],
[197, 184, 55],
...,
[186, 104, 46],
[254, 67, 40],
[ 96, 180, 188]],
[[ 1, 220, 190],
[ 47, 249, 174],
[ 85, 242, 42],
...,
[167, 32, 11],
[141, 140, 131],
[249, 9, 151]],
[[ 38, 178, 10],
[255, 188, 251],
[106, 62, 216],
...,
[242, 12, 226],
[227, 133, 94],
[ 13, 98, 189]],
...,
...
[118, 197, 154],
...,
[152, 113, 163],
[ 69, 135, 175],
[104, 193, 228]]]], dtype=uint16)
This numpy array is flattened with the numpy flatten method and the data shape set for the dim0 and dim1 parameters. The table is used for the inference input.
img_flattened = image.flatten()
table_with_flattened_array = pa.Table.from_pydict({"image": [img_flattened],
"dim0": [image.shape[0]],
"dim1": [image.shape[1]]},
input_schema
)
display(table_with_flattened_array)
pyarrow.Table
image: list<item: uint16>
child 0, item: uint16
dim0: int64
dim1: int64
----
image: [[[26,97,49,125,62,...,77,18,176,143,225]]]
dim0: [[6144]]
dim1: [[4096]]
The following example is the Wallaroo Custom Model 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:
flattened_image = input_data['image']
dim0 = input_data["dim0"][0]
dim1 = input_data["dim1"][0]
shape = (dim0, dim1)
image = input_data["image"].reshape(shape)
# //do some processing
# return the original flattened image and shape,
# with the virtual stain
return {
"image": flattened_image,
"dim0": dim0,
"dim1": dim1,
"virtual_stain": virtual_stain.flatten().astype(np.uint8)
}
Apache Arrow Table with Fixed Shape Tensor Arrays
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. In these cases, the data must be numpy arrays which are converted to fixed shape tensors.
The following example demonstrates a Wallaroo Custom Model Framework code to support this data scheme. Note that both the input and output schemas use the fix shape tensor data type.
The input and output schemas in Apache pyarrow format are used at model upload to specify the accepted data schemas.
input_schema = pa.schema([
pa.field('image', pa.fixed_shape_tensor(pa.uint16(), [6144, 4096, 3])),
])
# output schema
output_schema = pa.schema([
pa.field('image', pa.fixed_shape_tensor(pa.uint16(), [6144, 4096, 3])),
pa.field('virtual_stain', pa.fixed_shape_tensor(pa.uint8(), [6144, 4096, 3]))
])
The image data takes the image and converts it to an multi-dimensional numpy array, with each pixel’s RGB values output.
display(image)
array([[[[ 27, 159, 139],
[114, 107, 101],
[197, 184, 55],
...,
[186, 104, 46],
[254, 67, 40],
[ 96, 180, 188]],
[[ 1, 220, 190],
[ 47, 249, 174],
[ 85, 242, 42],
...,
[167, 32, 11],
[141, 140, 131],
[249, 9, 151]],
[[ 38, 178, 10],
[255, 188, 251],
[106, 62, 216],
...,
[242, 12, 226],
[227, 133, 94],
[ 13, 98, 189]],
...,
...
[118, 197, 154],
...,
[152, 113, 163],
[ 69, 135, 175],
[104, 193, 228]]]], dtype=uint16)
This numpy array is stored as an Apache Arrow table. Note that the data shape is preserved with its conversion of the numpy array to a an Apache Arrow fixed shape tensor array. This is used for the inference input.
In this example, the shape is updated by adding (1,) + image.shape to support batch processing a multiple images; these can be put into the same array to improve processing performance.
image.shape = (1,) + image.shape
table_with_fixed_shape_tensor_array = pa.Table.from_pydict({ "image": pa.FixedShapeTensorArray.from_numpy_ndarray(image)}, input_schema)
display(table_with_fixed_shape_tensor_array)
pyarrow.Table
image: extension<arrow.fixed_shape_tensor[value_type=uint16, shape=[6144,4096,3]]>
----
image: [[[27,159,139,114,107,...,135,175,104,193,228]]]
The following example is the Wallaroo Custom Model 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:
# //do some processing
# return the original shape and the virtual stain
return {
"image": input_data['image'].astype(np.uint16),
"virtual_stain": image_3d.astype(np.uint8)
}