39 lines
1.2 KiB
Python
39 lines
1.2 KiB
Python
import numpy as np
|
|
|
|
from mlflow.models import ModelSignature
|
|
from mlflow.types import Schema, TensorSpec
|
|
|
|
|
|
def get_model_signature(model):
|
|
def replace_none_in_shape(shape):
|
|
return [-1 if dim_size is None else dim_size for dim_size in shape]
|
|
|
|
input_shape = model.input_shape
|
|
input_dtype = model.input_dtype
|
|
output_shape = model.output_shape
|
|
output_dtype = model.compute_dtype
|
|
|
|
if isinstance(input_shape, list):
|
|
input_schema = Schema(
|
|
[
|
|
TensorSpec(np.dtype(input_dtype), replace_none_in_shape(shape))
|
|
for shape in input_shape
|
|
]
|
|
)
|
|
else:
|
|
input_schema = Schema(
|
|
[TensorSpec(np.dtype(input_dtype), replace_none_in_shape(input_shape))]
|
|
)
|
|
if isinstance(output_shape, list):
|
|
output_schema = Schema(
|
|
[
|
|
TensorSpec(np.dtype(output_dtype), replace_none_in_shape(shape))
|
|
for shape in output_shape
|
|
]
|
|
)
|
|
else:
|
|
output_schema = Schema(
|
|
[TensorSpec(np.dtype(output_dtype), replace_none_in_shape(output_shape))]
|
|
)
|
|
return ModelSignature(inputs=input_schema, outputs=output_schema)
|