trainY = tf.keras.utils.to_categorical(train_Y)
testY = tf.keras.utils.to_categorical(test_Y)
model = tf.keras.models.Sequential()
model.add(Conv2D(32, (3, 3), activation='relu', kernel_initializer='he_uniform', input_shape=(28, 28, 1)))
model.add(MaxPooling2D((2, 2)))
model.add(Flatten())
model.add(Dense(100, activation='relu', kernel_initializer='he_uniform'))
model.add(Dense(10, activation='softmax'))
opt = SGD(lr=0.01, momentum=0.9)
model.compile(optimizer=opt, loss='categorical_crossentropy', metrics=['accuracy'])
model.fit(trainX, trainY, epochs=1, batch_size=32, validation_data=(testX, testY))
input_schema = Schema([
TensorSpec(np.dtype(np.uint8), (-1, 28, 28, 1)),
])
output_schema = Schema([TensorSpec(np.dtype(np.float32), (-1, 10))])
signature = ModelSignature(inputs=input_schema, outputs=output_schema)
input_example = np.array([
[[ 0, 0, 0, 0],
[ 0, 134, 25, 56],
[253, 242, 195, 6],
[ 0, 93, 82, 82]],
[[ 0, 23, 46, 0],
[ 33, 13, 36, 166],
[ 76, 75, 0, 255],
[ 33, 44, 11, 82]]
], dtype=np.uint8)
mlflow.keras.log_model(model, "mnist_cnn", signature=signature, input_example=input_example)
def test_model_signature_with_tensorspec():
signature1 = ModelSignature(
inputs=Schema([
TensorSpec(np.dtype("float"), (-1, 28, 28)),
TensorSpec(np.dtype("int"), (-1, 10))
]),
outputs=Schema([TensorSpec(np.dtype("float"), (-1, 10))]),
)
signature2 = ModelSignature(
inputs=Schema([
TensorSpec(np.dtype("float"), (-1, 28, 28)),
TensorSpec(np.dtype("int"), (-1, 10))
]),
outputs=Schema([TensorSpec(np.dtype("float"), (-1, 10))]),
)
# Single type mismatch
assert signature1 == signature2
signature3 = ModelSignature(
inputs=Schema([
TensorSpec(np.dtype("float"), (-1, 28, 28)),
TensorSpec(np.dtype("int"), (-1, 10))
]),
outputs=Schema([TensorSpec(np.dtype("int"), (-1, 10))]),
)
assert signature3 != signature1
# Name mismatch
signature4 = ModelSignature(
inputs=Schema([
TensorSpec(np.dtype("float"), (-1, 28, 28)),
TensorSpec(np.dtype("int"), (-1, 10))
]),
outputs=Schema([TensorSpec(np.dtype("float"), (-1, 10), "misMatch")]),
)
assert signature3 != signature4
as_json = json.dumps(signature1.to_dict())
signature5 = ModelSignature.from_dict(json.loads(as_json))
assert signature1 == signature5
# Test with name
signature6 = ModelSignature(
inputs=Schema([
TensorSpec(np.dtype("float"), (-1, 28, 28), name="image"),
TensorSpec(np.dtype("int"), (-1, 10), name="metadata"),
]),
outputs=Schema(
[TensorSpec(np.dtype("float"), (-1, 10), name="outputs")]),
)
signature7 = ModelSignature(
inputs=Schema([
TensorSpec(np.dtype("float"), (-1, 28, 28), name="image"),
TensorSpec(np.dtype("int"), (-1, 10), name="metadata"),
]),
outputs=Schema(
[TensorSpec(np.dtype("float"), (-1, 10), name="outputs")]),
)
assert signature6 == signature7
assert signature1 != signature6
# Test w/o output
signature8 = ModelSignature(inputs=Schema(
[TensorSpec(np.dtype("float"), (-1, 28, 28))]),
outputs=None)
as_json = json.dumps(signature8.to_dict())
signature9 = ModelSignature.from_dict(json.loads(as_json))
assert signature8 == signature9
def test_model_save_load():
m = Model(
artifact_path="some/path",
run_id="123",
flavors={
"flavor1": {
"a": 1,
"b": 2
},
"flavor2": {
"x": 1,
"y": 2
}
},
signature=ModelSignature(
inputs=Schema([ColSpec("integer", "x"),
ColSpec("integer", "y")]),
outputs=Schema([ColSpec(name=None, type="double")]),
),
saved_input_example_info={
"x": 1,
"y": 2
},
)
assert m.get_input_schema() == m.signature.inputs
assert m.get_output_schema() == m.signature.outputs
x = Model(artifact_path="some/other/path", run_id="1234")
assert x.get_input_schema() is None
assert x.get_output_schema() is None
n = Model(
artifact_path="some/path",
run_id="123",
flavors={
"flavor1": {
"a": 1,
"b": 2
},
"flavor2": {
"x": 1,
"y": 2
}
},
signature=ModelSignature(
inputs=Schema([ColSpec("integer", "x"),
ColSpec("integer", "y")]),
outputs=Schema([ColSpec(name=None, type="double")]),
),
saved_input_example_info={
"x": 1,
"y": 2
},
)
n.utc_time_created = m.utc_time_created
assert m == n
n.signature = None
assert m != n
with TempDir() as tmp:
m.save(tmp.path("model"))
o = Model.load(tmp.path("model"))
assert m == o
assert m.to_json() == o.to_json()
assert m.to_yaml() == o.to_yaml()
def _infer_signature(onnx_model):
onnx_model_bytes = onnx_model.SerializeToString()
onnx_runtime = onnxruntime.InferenceSession(onnx_model_bytes)
inputs = _infer_schema(onnx_runtime.get_inputs())
outputs = _infer_schema(onnx_runtime.get_outputs())
return ModelSignature(inputs, outputs)