Inserting a TensorFlow network into a Nengo model¶
TensorFlow comes with a wide range of pre-defined deep learning models, which we might want to incorporate into a Nengo model. For example, suppose we are building a biological reinforcement learning model, but we’d like the inputs to our model to be natural images rather than artificial vectors. We could load a vision network from TensorFlow, insert it into our model using NengoDL, and then build the rest of our model using normal Nengo syntax.
In this example we’ll show how to use TensorNodes to insert a pre-trained TensorFlow model (Inception-v1) into Nengo.
In [ ]:
%matplotlib inline import sys import os from urllib.request import urlopen import io import shutil import stat import numpy as np import matplotlib.pyplot as plt from PIL import Image import tensorflow as tf import tensorflow.contrib.slim as slim; import nengo import nengo_dl
TensorFlow provides a number of pre-defined models in the tensorflow/models repository. These are not included when you install TensorFlow, so we need to separately clone that repository and import the components we need.
In [ ]:
!git clone https://github.com/tensorflow/models sys.path.append(os.path.join(".", "models", "slim")) from datasets import dataset_utils, imagenet from nets import inception from preprocessing import inception_preprocessing
We will use a
insert our TensorFlow code into Nengo.
very similarly to
nengo.Node, except instead of using the node to
insert Python code into our model we will use it to insert TensorFlow
The first thing we need to do is define our TensorNode output. This
should be a function that accepts the current simulation time (and,
optionally, a batch of vectors) as input, and produces a batch of
vectors as output. All of these variables will be represented as
tf.Tensor objects, and the internal operations of the TensorNode
will be implemented with TensorFlow operations. For example, we could
use a TensorNode to output a
In [ ]:
with nengo.Network() as net: node = nengo_dl.TensorNode(lambda t: tf.reshape(tf.sin(t), (1, 1))) p = nengo.Probe(node) with nengo_dl.Simulator(net) as sim: sim.run(5.0) plt.figure() plt.plot(sim.trange(), sim.data[p])
However, outputing a
sin function is something we could do more
easily with a regular
nengo.Node. The main use case for
nengo_dl.TensorNode is to work with artificial neural networks that
are not easily defined in Nengo.
In this case we’re going to build a TensorNode that encapsulates the Inception-v1 network. Inception-v1 isn’t state-of-the-art anymore (we’re up to Inception-v4 now), but it is relatively small so it will be quick to download/run in this example. However, this same approach could be used for any TensorFlow network.
Inception-v1 performs image classification; if we show it an image, it will output a set of probabilities for the 1000 different object types it is trained to classify. So if we show it an image of a tree it should output a high probability for the “tree” class and a low probability for the “car” class.
The first thing we’ll do is download a sample image to test our network with (you could use a different image if you want).
In [ ]:
url = 'https://upload.wikimedia.org/wikipedia/commons/7/70/EnglishCockerSpaniel_simon.jpg' image_string = urlopen(url).read() image = np.array(Image.open(io.BytesIO(image_string))) image_shape = image.shape # display the test image plt.figure() plt.imshow(image) plt.axis('off');
Now we’re ready to create our TensorNode. Instead of using a function
for our TensorNode output, in this case we’ll use a callable class. This
allows us to include a
pre_build function in our TensorNode. NengoDL
will call this function once when the model is first constructed, so we
can use this function to perform any initial setup required for our
In [ ]:
checkpoints_dir = '/tmp/checkpoints' class InceptionNode(object): def pre_build(self, *args): # the shape of the inputs to the inception network self.input_shape = inception.inception_v1.default_image_size # download model checkpoint file if not tf.gfile.Exists(checkpoints_dir): tf.gfile.MakeDirs(checkpoints_dir) url = "http://download.tensorflow.org/models/inception_v1_2016_08_28.tar.gz" dataset_utils.download_and_uncompress_tarball(url, checkpoints_dir) def __call__(self, t, x): # this is the function that will be executed each timestep while the # network is running # convert our input vector to the shape/dtype of the input image image = tf.reshape(tf.cast(x, tf.uint8), image_shape) # reshape the image to the shape expected by the inception network processed_image = inception_preprocessing.preprocess_image( image, self.input_shape, self.input_shape, is_training=False) processed_images = tf.expand_dims(processed_image, 0) # create inception network with slim.arg_scope(inception.inception_v1_arg_scope()): logits, _ = inception.inception_v1(processed_images, num_classes=1001, is_training=False) probabilities = tf.nn.softmax(logits) # return our classification probabilites return probabilities
Next we create a Nengo Network, containing our TensorNode.
In [ ]:
with nengo.Network() as net: # create a normal input node to feed in our test image input_node = nengo.Node(output=image.flatten()) # create our TensorNode containing the InceptionNode() we defined # above. we also need to specify size_in (the dimensionality of # our input vectors, the flattened images) and size_out (the number # of classification classes output by the inception network) incep_node = nengo_dl.TensorNode( InceptionNode(), size_in=np.prod(image_shape), size_out=1001) # connect up our input to our inception node nengo.Connection(input_node, incep_node, synapse=None) # add some probes to collect data input_p = nengo.Probe(input_node) incep_p = nengo.Probe(incep_node)
Note that at this point we could connect up the output of
to any other part of our network, if this was part of a larger model.
But to keep this example simple we’ll stop here.
Next we’ll load the pre-trained weights for the Inception network. If we
wanted we could train the network from scratch using the
function, but that would take a long time and require some expertise in
training deep networks.
In [ ]:
sim = nengo_dl.Simulator(net) with sim.tensor_graph.graph.as_default(): # load checkpoint file init_fn = slim.assign_from_checkpoint_fn( os.path.join(checkpoints_dir, 'inception_v1.ckpt'), slim.get_model_variables('InceptionV1')) # update the parameters of the model init_fn(sim.sess)
And finally, all that’s left is to run our network, using our example image as input, and check the output
In [ ]:
# run the network for one timestep sim.step() # sort the output labels based on the classification probabilites # output from the network probabilities = sim.data[incep_p] sorted_inds = [i for i in sorted(enumerate(-probabilities), key=lambda x: x)] # print top 5 classes names = imagenet.create_readable_names_for_imagenet_labels() for i in range(5): index = sorted_inds[i] print('Probability %0.2f%% => [%s]' % ( probabilities[index] * 100, names[index])) # display the test image plt.figure() plt.imshow(sim.data[input_p].reshape(image_shape).astype(np.uint8)) plt.axis('off');
In [ ]:
In [ ]:
# delete the models repo we cloned def onerror(func, path, exc_info): if not os.access(path, os.W_OK): os.chmod(path, stat.S_IWUSR) func(path) else: raise shutil.rmtree("models", onerror=onerror)