TensorFlow
TensorFlow is a supported framework on Truss. To package a TensorFlow model, follow the steps below or run this Google Colab notebook.
If you're using a Jupyter notebook, add a line to install the
tensorflow and truss packages. Otherwise, ensure the packages are installed in your Python environment.!pip install --upgrade tensorflow truss
# For help installing tensorflow, see https://www.tensorflow.org/install/pip
Truss officially supports
tensorflow version 2.4.0 or higher. Especially if you're using an online notebook environment like Google Colab or a bundle of packages like Anaconda, ensure that the version you are using is supported. If it's not, use the --upgrade flag and pip will install the most recent version.This is the part you want to replace with your own code. Using TensorFlow, build a machine learning model and keep it in-memory.
import tensorflow as tf
#Creates tensorflow model
model = tf.keras.applications.ResNet50V2(
include_top=True,
weights="imagenet",
classifier_activation="softmax",
)
Use the
create command to package your model into a Truss.from truss import create
tr = create(model, target_directory="tensorflow_truss")
Check the target directory to see your new Truss!
In your newly created Truss, open
model/model.py and add pre- and post-processing functions as follows.First, add the following imports at the top of the file:
import requests
import tempfile
import tensorflow as tf
from scipy.special import softmax
Then, update the pre-processing function to:
def preprocess(self, model_input: Any) -> Any:
"""Preprocess step for ResNet"""
request = requests.get(model_input)
with tempfile.NamedTemporaryFile() as f:
f.write(request.content)
f.seek(0)
input_image = tf.image.decode_png(tf.io.read_file(f.name))
preprocessed_image = tf.keras.applications.resnet_v2.preprocess_input(
tf.image.resize([input_image], (224, 224))
)
return np.array(preprocessed_image)
Finally, update the post-processing function to:
def postprocess(self, model_output: Dict) -> Dict:
"""Post process step for ResNet"""
class_predictions = model_output["predictions"][0]
LABELS = requests.get(
"https://storage.googleapis.com/download.tensorflow.org/data/ImageNetLabels.txt"
).text.split("\n")
class_probabilities = softmax(class_predictions)
top_probability_indices = class_probabilities.argsort()[::-1][:5].tolist()
return {
LABELS[index]: 100 * class_probabilities[index].round(3)
for index in top_probability_indices
}
With these functions in place, you can invoke the model and pass it a URL, as in:
from truss import load
tr = load("tensorflow_truss")
tr.predict("https://github.com/pytorch/hub/raw/master/images/dog.jpg")
Last modified 18d ago