# 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 2mo ago