TensorFlow tf.py_func(): Run Python Function in TensorFlow Graph

By | December 31, 2019
0 Comment

TensorFlow tf.py_func() can allow us to run python script in tensorflow graph. In this tutorial, we will write some examples to show you how to use it correctly.

Syntax

tf.py_func(
    func,
    inp,
    Tout,
    stateful=True,
    name=None
)

Wraps a python function and uses it as a TensorFlow op.

Parameters explained

func: a python function you plan to run in tensorflow graph, it can be written by python, numpy etc.

inp: a list of tensors, these tensors will be passed into func to execute and compute, they are parameters of func. They will be converted to numpy.ndarray in func.

Tout: a list of tensor type, which tells what types of data func will return.

Here is an example to show how to use it.

import tensorflow as tf
import numpy as np

def computeX(x):
    print(type(x))
    return x+x, np.multiply(x, x), np.exp(x)

r = tf.py_func(computeX, [w1], [tf.float32, tf.float32, tf.float32])

nit = tf.global_variables_initializer() 
init_local = tf.local_variables_initializer()

with tf.Session() as sess:
    sess.run([init, init_local])
    print(sess.run([r]))

Example exaplained

r = tf.py_func(computeX, [w1], [tf.float32, tf.float32, tf.float32])

This code means:

We have created a python function computeX, it will use one parameter and it will return three value, the type of these three value is float, float and float.

[w1]: this tensor will be passed into computeX() and it will be converted into a numpy.ndarray

[tf.float32, tf.float32, tf.float32]: it means computeX() will return three value, the type of these three is float, float and float.

As to code:

def computeX(x):
    print(type(x))
    return x+x, np.multiply(x, x), np.exp(x)

In computeX(x), it will receive a parameter, which is a tensor. This parameter x will be converted to numpy.ndarray.

Run this python script, you will get result:

<class 'numpy.ndarray'>
[[array([[ 2.,  4.,  6.,  8., 10.],
       [12., 14., 16., 18.,  0.]], dtype=float32), array([[ 1.,  4.,  9., 16., 25.],
       [36., 49., 64., 81.,  0.]], dtype=float32), array([[2.7182817e+00, 7.3890562e+00, 2.0085537e+01, 5.4598148e+01,
        1.4841316e+02],
       [4.0342880e+02, 1.0966332e+03, 2.9809580e+03, 8.1030840e+03,
        1.0000000e+00]], dtype=float32)]]

From the result you will find:

x in computeX(x) is converted to numpy.ndarray

computeX(x) will return three values.

Notice:

1. If you use python function to compute a tensor in tf.py_func(). The gradient will be None.

Here is an example:

g = tf.gradients(r, w1)
init = tf.global_variables_initializer() 
init_local = tf.local_variables_initializer()

with tf.Session() as sess:
    sess.run([init, init_local])
    print(sess.run([g]))

Run this code, you will get error:

TypeError: Fetch argument None has invalid type <class ‘NoneType’>

2.The body of the function (computeX) will not be serialized in a GraphDef. Therefore, you should not use tf.py_func() if you need to serialize your model and restore it in a different environment.

 

Leave a Reply