Tensorflow¶

Users can now use Tensorflow (v1 and v2 supported) with BentoML with the following API: load, save, and load_runner as follow:

import bentoml
import tensorflow as tf

class NativeModel(tf.Module):
    def __init__(self):
    super().__init__()
    self.weights = np.asfarray([[1.0], [1.0], [1.0], [1.0], [1.0]])
    self.dense = lambda inputs: tf.matmul(inputs, self.weights)

    @tf.function(
    input_signature=[tf.TensorSpec(shape=[1, 5], dtype=tf.float64, name="inputs")]
    )
    def __call__(self, inputs):
    return self.dense(inputs)

model = NativeModel()

# `save` a given model and retrieve coresponding tag:
tag = bentoml.tensorflow.save("native_tf_module", model)

# retrieve metadata with `bentoml.models.get`:
metadata = bentoml.models.get(tag)

# `load` the model back in memory:
model = bentoml.tensorflow.load("cancer_clf:latest")

# Run a given model under `Runner` abstraction with `load_runner`
_data = [[1.1, 2.2]]
_tensor = tf.constant(_data)
runner = bentoml.tensorflow.load_runner("native_tf_module:latest"")
runner.run(_tensor)

We also offer import_from_tfhub which enables users to import model from Tensorflow Hub and use it with BentoML:

import tensorflow_text as text
import bentoml

tag = bentoml.tensorflow.import_from_tfhub("https://tfhub.dev/tensorflow/bert_en_uncased_preprocess/3")

# load model back with `load`:
model = bentoml.tensorflow.load(tag, load_as_wrapper=True)

Note

You can find more examples for Tensorflow in our gallery repo.

bentoml.tensorflow.save(name, model, *, signatures=None, options=None, labels=None, custom_objects=None, metadata=None)¶

Save a model instance to BentoML modelstore.

Parameters

name (str) – Name for given model instance. This should pass Python identifier check.
model (Union[keras.Model, tf.Module, path-like objects]) – Instance of model to be saved
labels (Dict[str, str], optional, default to None) – user-defined labels for managing models, e.g. team=nlp, stage=dev
custom_objects (Dict[str, Any]], optional, default to None) – user-defined additional python objects to be saved alongside the model, e.g. a tokenizer instance, preprocessor function, model configuration json
metadata (Dict[str, Any], optional, default to None) – Custom metadata for given model.
model_store (ModelStore, default to BentoMLContainer.model_store) – BentoML modelstore, provided by DI Container.
signatures (Union[Callable[..., Any], dict], optional, default to None) – Refers to Signatures explanation from Tensorflow documentation for more information.
options (tf.saved_model.SaveOptions, optional, default to None) – tf.saved_model.SaveOptions object that specifies options for saving.

Raises

ValueError – If obj is not trackable.

Returns

A tag with a format name:version where name is the user-defined model’s name, and a generated version by BentoML.

Return type

Tag

Examples:

import tensorflow as tf
import numpy as np
import bentoml

class NativeModel(tf.Module):
    def __init__(self):
        super().__init__()
        self.weights = np.asfarray([[1.0], [1.0], [1.0], [1.0], [1.0]])
        self.dense = lambda inputs: tf.matmul(inputs, self.weights)

    @tf.function(
        input_signature=[tf.TensorSpec(shape=[1, 5], dtype=tf.float64, name="inputs")]
    )
    def __call__(self, inputs):
        return self.dense(inputs)

# then save the given model to BentoML modelstore:
model = NativeModel()
tag = bentoml.tensorflow.save("native_toy", model)

Note

bentoml.tensorflow.save API also support saving RaggedTensor model and Keras model. If you choose to save a Keras model with bentoml.tensorflow.save, then the model will be saved under a SavedModel format instead of h5.

bentoml.tensorflow.load(bento_tag, tags=None, options=None, load_as_hub_module=None, model_store=<simple_di.providers.SingletonFactory object>)¶

Load a model from BentoML local modelstore with given name.

Parameters

bento_tag (Union[str, Tag]) – Tag of a saved model in BentoML local modelstore.
tags (str, optional, defaults to None) – A set of strings specifying the graph variant to use, if loading from a v1 module.
options (tensorflow.saved_model.SaveOptions, optional, default to None) – tensorflow.saved_model.LoadOptions object that specifies options for loading. This argument can only be used from TensorFlow 2.3 onwards.
load_as_hub_module (bool, optional, default to True) – Load the given weight that is saved from tfhub as either hub.KerasLayer or hub.Module. The latter only applies for TF1.
model_store (ModelStore, default to BentoMLContainer.model_store) – BentoML modelstore, provided by DI Container.

Returns

an instance of SavedModel format from BentoML modelstore.

Return type

SavedModel

Examples:

import bentoml

# load a model back into memory
model = bentoml.tensorflow.load("my_tensorflow_model")

bentoml.tensorflow.load_runner(tag, *, predict_fn_name='__call__', device_id='CPU:0', name=None, partial_kwargs=None)¶

Runner represents a unit of serving logic that can be scaled horizontally to maximize throughput. bentoml.tensorflow.load_runner implements a Runner class that wrap around a Tensorflow model, which optimize it for the BentoML runtime.

Parameters

tag (Union[str, Tag]) – Tag of a saved model in BentoML local modelstore.
predict_fn_name (str, default to __call__) – Inference function to be used.
partial_kwargs (Dict[str, Any], optional, default to None) – Dictionary of partial kwargs that can be shared across different model.
device_id (str, optional, default to the first CPU) – Optional devices to put the given model on. Refers to Logical Devices from TF documentation.

Returns

Runner instances for bentoml.tensorflow model

Return type

Runner

Examples:

import bentoml

# load a runner from a given flag
runner = bentoml.tensorflow.load_runner(tag)

# load a runner on GPU:0
runner = bentoml.tensorflow.load_runner(tag, resource_quota=dict(gpus=0), device_id="GPU:0")

bentoml.tensorflow.import_from_tfhub(identifier, name=None, labels=None, custom_objects=None, metadata=None)¶

Import a model from Tensorflow Hub to BentoML modelstore.

Parameters

identifier (Union[str, tensorflow_hub.Module, tensorflow_hub.KerasLayer]) –
Identifier accepts two type of inputs:
- if type of identifier either of type tensorflow_hub.Module (legacy tensorflow_hub) or tensorflow_hub.KerasLayer (tensorflow_hub), then we will save the given model to a SavedModel format.
- if type of identifier is a str, we assume that this is the URI retrieved from Tensorflow Hub. We then clean the given URI, and get a local copy of a given model to BentoML modelstore. name (str, optional, defaults to None): An optional name for the model. If identifier is a str, then name can be autogenerated from the given URI.
name (str, optional, default to None) – Optional name for the saved model. If None, then name will be generated from identifier.
labels (Dict[str, str], optional, default to None) – user-defined labels for managing models, e.g. team=nlp, stage=dev
custom_objects (Dict[str, Any]], optional, default to None) – user-defined additional python objects to be saved alongside the model, e.g. a tokenizer instance, preprocessor function, model configuration json
metadata (Dict[str, Any], optional, default to None) – Custom metadata for given model.
model_store (ModelStore, default to BentoMLContainer.model_store) – BentoML modelstore, provided by DI Container.

Returns

A Tag object that can be used to retrieve the model with bentoml.tensorflow.load():

Return type

Tag

Example for importing a model from Tensorflow Hub:

import tensorflow_text as text # noqa # pylint: disable
import bentoml

tag = bentoml.tensorflow.import_from_tfhub("https://tfhub.dev/tensorflow/bert_en_uncased_preprocess/3")

# load model back with `load`:
model = bentoml.tensorflow.load(tag, load_as_hub_module=True)

Example for importing a custom Tensorflow Hub model:

import tensorflow as tf
import tensorflow_hub as hub
import bentoml

def _plus_one_model_tf2():
    obj = tf.train.Checkpoint()

    @tf.function(input_signature=[tf.TensorSpec(None, dtype=tf.float32)])
    def plus_one(x):
        return x + 1

    obj.__call__ = plus_one
    return obj

# then save the given model to BentoML modelstore:
model = _plus_one_model_tf2()
tag = bentoml.tensorflow.import_from_tfhub(model)