I am trying to write some custom TensorFlow functions in python (using tf.py_func) where I want to calculate both the results and the gradients in python. I'm using the gradient_override_map trick (for example from from https://gist.github.com/harpone/3453185b41d8d985356cbe5e57d67342 and How to make a custom activation function with only Python in Tensorflow?).
However, while the function in the forward direction gets a numpy array as an input, the function for the gradient gets Tensors. This is a problem, depending on when the function gets called, because there may not be a default session, and/or there may not be a feed_dict with all the required values yet (for example, in a tf.train optimizer).
How do I do a py_func where both the forward and backward functions get (and return) numpy arrays?
Sample code:
import tensorflow as tf
import numpy as np
import matplotlib.pyplot as plt
def sin_func(x):
return np.sin(x)
def sin_grad_func(op, grad):
x = op.inputs[0].eval()
grad = grad.eval() # <--- this is what I'd like to avoid
output_grad = np.cos(x) * grad
return tf.convert_to_tensor(output_grad)
def py_func(func, inp, Tout, stateful=True, name=None, grad_func=None):
grad_name = 'PyFuncGrad_' + str(np.random.randint(0, 1E+8))
tf.RegisterGradient(grad_name)(grad_func)
g = tf.get_default_graph()
with g.gradient_override_map({"PyFunc": grad_name}):
return tf.py_func(func, inp, Tout, stateful=stateful, name=name)
with tf.Session() as sess:
np_x = np.linspace(0, np.pi, num=1000, dtype=np.float32)
x = tf.constant(np_x)
y = py_func(sin_func,
[x],
[tf.float32],
name='np_sin',
grad_func=sin_grad_func)
y = y[0]
gr = tf.gradients(y, [x])
tf.global_variables_initializer().run()
plt.plot(y.eval())
plt.plot(gr[0].eval())
If you want to include arbitrary Python code in your gradient function, the easiest solution is to create another tf.py_func() inside sin_grad_func():
def sin_grad_func_impl(x, grad):
return np.cos(x) * grad
def sin_grad_func(op, grad):
return tf.py_func(sin_grad_func_impl, [x, grad], grad.dtype)
Related
Why is GradientTape returning None when I use numpy math
I am trying to understand tensorflow GradientTape calculation for RL loss function. When I call a function using np.math the GradientTape returns None. If I use tf.math in the function it works fine. I have looked at tf-agents like ppo and sac and they are doing exactly(?) what I am trying to do (I have tried at last 50 other versions).
What's wrong in the code below? What am I missing?
window 10, python 3.6.8, tensorflow 2.0.0
ref:https://github.com/chagmgang/tf2.0_reinforcement_learning/blob/master/policy/ppo.py
import numpy as np
import tensorflow as tf
import matplotlib.pyplot as plt
def my_loss1(x):
y=tf.sin(x)
y=tf.abs(y)
return y
def my_loss2(x):
y=np.sin(x)
y=np.abs(y)
return y
def main(ver):
x = np.linspace(0,10,25)
dsin_dx=np.cos(x)
xx = tf.constant(x)
with tf.GradientTape() as tape:
tape.watch(xx)
if ver==0:
# my_loss1 with tf math
loss1=my_loss1(xx)
if ver==1:
#my loss with numpy math
loss1=my_loss2(np.array(xx))
loss1 = tf.convert_to_tensor(loss1, dtype=tf.float64)
print(loss1)
loss=tf.reduce_sum(loss1)
print('loss=',loss)
grads = tape.gradient(loss, xx)
fig, ax = plt.subplots(2)
ax[0].plot(x,loss1,'r')
print('grads', grads)
if not grads is None:
ax[1].plot(x, grads)
ax[1].plot(x,dsin_dx)
plt.show()
if __name__ == '__main__':
main(ver=0) # This works ok
main(ver=1) # This returns grads = None
The problem is that the Gradient tape only records tensors. Numpy variables are not recorded why the gradient can't be calqulated in case ver=1. Loss1 in ver1 looks identical to loss1 in ver=0 but the dependentsy to xx is broken by numpy.
My ref. has this error when calculation get_gaes() and the calculation of the grads is incorrect.
This layer is static, it is a pseudo function. In the forward propagation it doesn't do anything (identity function). In the back propagation however, it multiplies the gradient by -1. There are lots of implementations on github but they don't work with TF 2.0.
Here's one for reference.
import tensorflow as tf
from tensorflow.python.framework import ops
class FlipGradientBuilder(object):
def __init__(self):
self.num_calls = 0
def __call__(self, x, l=1.0):
grad_name = "FlipGradient%d" % self.num_calls
#ops.RegisterGradient(grad_name)
def _flip_gradients(op, grad):
return [tf.negative(grad) * l]
g = tf.get_default_graph()
with g.gradient_override_map({"Identity": grad_name}):
y = tf.identity(x)
self.num_calls += 1
return y
flip_gradient = FlipGradientBuilder()
Dummy op that reverses the gradients
This can be done using the decorator tf.custom_gradient, as described in this example:
#tf.custom_gradient
def grad_reverse(x):
y = tf.identity(x)
def custom_grad(dy):
return -dy
return y, custom_grad
Then, you can just use it as if it is a normal TensorFlow op, for example:
z = encoder(x)
r = grad_reverse(z)
y = decoder(r)
Keras API?
A great convenience of TF 2.0 is it's native support for Keras API. You can define a custom GradReverse op and enjoy the convenience of Keras:
class GradReverse(tf.keras.layers.Layer):
def __init__(self):
super().__init__()
def call(self, x):
return grad_reverse(x)
Then, you can use this layer as any other layers of Keras, for example:
model = Sequential()
conv = tf.keras.layers.Conv2D(...)(inp)
cust = CustomLayer()(conv)
flat = tf.keras.layers.Flatten()(cust)
fc = tf.keras.layers.Dense(num_classes)(flat)
model = tf.keras.models.Model(inputs=[inp], outputs=[fc])
model.compile(loss=..., optimizer=...)
model.fit(...)
I am trying to develop a new model within gpflow. In order to debug it I need to know shapes and values of tensors during execution of the graph.
I tried the below based on printing tensor values in tensorflow, but nothing is printed to the console.
import numpy as np
import sys
import gpflow
from gpflow.mean_functions import MeanFunction
from gpflow.decors import params_as_tensors
class Log(MeanFunction):
"""
:math:`y_i = \log(x_i)`
"""
def __init__(self):
MeanFunction.__init__(self)
#params_as_tensors
def __call__(self, X):
# I want to figure out the shape of X here
tf.print(tf.shape(X), output_stream=sys.stdout)
# Returns the natural logarithm of the input
return tf.log(X)
# Test gpflow implementation
sess = tf.InteractiveSession()
with sess.as_default(), sess.graph.as_default():
X = np.random.uniform(size=[100, 1])
y = np.random.uniform(size=[100, 1])
m = gpflow.models.GPR(X=X, Y=y, mean_function=Log(), kern=gpflow.kernels.RBF(input_dim=1))
You're on the right track. According to the TensorFlow docs [1], you need to wrap tf.print() in a tf.control_dependencies() context manager to make sure it's run, when in graph model. GPflow currently works in graph model. GPflow 2.0, which is indevelopment, will allow usage in eager mode.
#params_as_tensors
def __call__(self, X):
# I want to figure out the shape of X here
print_op = tf.print(tf.shape(X), output_stream=sys.stdout)
with tf.control_dependencies([print_op]):
log_calc = tf.log(X)
# Returns the natural logarithm of the input
return log_calc
[1] https://www.tensorflow.org/api_docs/python/tf/print
I have a non trivial input pipeline that from_generator is perfect for...
dataset = tf.data.Dataset.from_generator(complex_img_label_generator,
(tf.int32, tf.string))
dataset = dataset.batch(64)
iter = dataset.make_one_shot_iterator()
imgs, labels = iter.get_next()
Where complex_img_label_generator dynamically generates images and returns a numpy array representing a (H, W, 3) image and a simple string label. The processing not something I can represent as reading from files and tf.image operations.
My question is about how to parallise the generator? How do I have N of these generators running in their own threads.
One thought was to use dataset.map with num_parallel_calls to handle the threading; but the map operates on tensors... Another thought was to create multiple generators each with it's own prefetch and somehow join them, but I can't see how I'd join N generator streams?
Any canonical examples I could follow?
Turns out I can use Dataset.map if I make the generator super lightweight (only generating meta data) and then move the actual heavy lighting into a stateless function. This way I can parallelise just the heavy lifting part with .map using a py_func.
Works; but feels a tad clumsy... Would be great to be able to just add num_parallel_calls to from_generator :)
def pure_numpy_and_pil_complex_calculation(metadata, label):
# some complex pil and numpy work nothing to do with tf
...
dataset = tf.data.Dataset.from_generator(lightweight_generator,
output_types=(tf.string, # metadata
tf.string)) # label
def wrapped_complex_calulation(metadata, label):
return tf.py_func(func = pure_numpy_and_pil_complex_calculation,
inp = (metadata, label),
Tout = (tf.uint8, # (H,W,3) img
tf.string)) # label
dataset = dataset.map(wrapped_complex_calulation,
num_parallel_calls=8)
dataset = dataset.batch(64)
iter = dataset.make_one_shot_iterator()
imgs, labels = iter.get_next()
I am working on a from_indexable for tf.data.Dataset https://github.com/tensorflow/tensorflow/issues/14448
The advantage for from_indexable is that it can be parallelized, while a python generator cannot be parallelized.
The function from_indexable makes a tf.data.range, wraps the indexable in a generalized tf.py_func and calls map.
For those that want now a from_indexable, here the lib code
import tensorflow as tf
import numpy as np
from tensorflow.python.framework import tensor_shape
from tensorflow.python.util import nest
def py_func_decorator(output_types=None, output_shapes=None, stateful=True, name=None):
def decorator(func):
def call(*args):
nonlocal output_shapes
flat_output_types = nest.flatten(output_types)
flat_values = tf.py_func(
func,
inp=args,
Tout=flat_output_types,
stateful=stateful, name=name
)
if output_shapes is not None:
# I am not sure if this is nessesary
output_shapes = nest.map_structure_up_to(
output_types, tensor_shape.as_shape, output_shapes)
flattened_shapes = nest.flatten_up_to(output_types, output_shapes)
for ret_t, shape in zip(flat_values, flattened_shapes):
ret_t.set_shape(shape)
return nest.pack_sequence_as(output_types, flat_values)
return call
return decorator
def from_indexable(iterator, output_types, output_shapes=None, num_parallel_calls=None, stateful=True, name=None):
ds = tf.data.Dataset.range(len(iterator))
#py_func_decorator(output_types, output_shapes, stateful=stateful, name=name)
def index_to_entry(index):
return iterator[index]
return ds.map(index_to_entry, num_parallel_calls=num_parallel_calls)
and here an example (Note: from_indexable has a num_parallel_calls argument)
class PyDataSet:
def __len__(self):
return 20
def __getitem__(self, item):
return np.random.normal(size=(item+1, 10))
ds = from_indexable(PyDataSet(), output_types=tf.float64, output_shapes=[None, 10])
it = ds.make_one_shot_iterator()
entry = it.get_next()
with tf.Session() as sess:
print(sess.run(entry).shape)
print(sess.run(entry).shape)
Update June 10, 2018:
Since https://github.com/tensorflow/tensorflow/pull/15121 is merged, the code for from_indexable simplifies to:
import tensorflow as tf
def py_func_decorator(output_types=None, output_shapes=None, stateful=True, name=None):
def decorator(func):
def call(*args, **kwargs):
return tf.contrib.framework.py_func(
func=func,
args=args, kwargs=kwargs,
output_types=output_types, output_shapes=output_shapes,
stateful=stateful, name=name
)
return call
return decorator
def from_indexable(iterator, output_types, output_shapes=None, num_parallel_calls=None, stateful=True, name=None):
ds = tf.data.Dataset.range(len(iterator))
#py_func_decorator(output_types, output_shapes, stateful=stateful, name=name)
def index_to_entry(index):
return iterator[index]
return ds.map(index_to_entry, num_parallel_calls=num_parallel_calls)
Limiting the work done in the generator to a minimum and parallelizing the expensive processing using a map is sensible.
Alternatively, you can "join" multiple generators using parallel_interleave as follows:
def generator(n):
# returns n-th generator function
def dataset(n):
return tf.data.Dataset.from_generator(generator(n))
ds = tf.data.Dataset.range(N).apply(tf.contrib.data.parallel_interleave(dataset, cycle_lenght=N))
# where N is the number of generators you use
In tensorflow, we can define our own op and its gradient by:
https://gist.github.com/harpone/3453185b41d8d985356cbe5e57d67342
However, can we modify any variable in the computational graph in these python functions. For example in the "_MySquareGrad" function?
I assume we can get the variable by:
var = tf.get_variable('var')
and then do something to change its value and then assign it back?
e.g.
tmp = var*10
var.assign(tmp)
Thanks!
Also when we do var*10, do we have to convert it to numpy?
Background: I'm familiar with automatic differentiation, but new to Tensorflow and Python. So please point out any syntactic problem and let me know if my intention is clear.
You can modify the variables in the computational graph in these python functions. Your example code with tmp = var*10 will work and does not convert anything to numpy.
In fact you should try to avoid converting to numpy as much as possible since it will slow down the computation.
edit:
You can include your code to the gradient computation graph of the _MySquareGrad function doing this:
def _MySquareGrad(op, grad):
#first get a Variable that was created using tf.get_variable()
with tf.variable_scope("", reuse=True):
var = tf.get_variable('var')
#now create the assign graph:
tmp = var*10.
assign_op = var.assign(tmp)
#now make the assign operation part of the grad calculation graph:
with tf.control_dependencies([assign_op]):
x = tf.identity(op.inputs[0])
return grad * 20 * x
Here is a working example:
import tensorflow as tf
from tensorflow.python.framework import ops
import numpy as np
# Define custom py_func which takes also a grad op as argument:
def py_func(func, inp, Tout, stateful=True, name=None, grad=None):
# Need to generate a unique name to avoid duplicates:
rnd_name = 'PyFuncGrad' + str(np.random.randint(0, 1E+8))
tf.RegisterGradient(rnd_name)(grad) # see _MySquareGrad for grad example
g = tf.get_default_graph()
with g.gradient_override_map({"PyFunc": rnd_name}):
return tf.py_func(func, inp, Tout, stateful=stateful, name=name)
# Def custom square function using np.square instead of tf.square:
def mysquare(x, name=None):
with ops.name_scope(name, "Mysquare", [x]) as name:
sqr_x = py_func(np.square,
[x],
[tf.float32],
name=name,
grad=_MySquareGrad) # <-- here's the call to the gradient
return sqr_x[0]
### Actual gradient:
##def _MySquareGrad(op, grad):
##x = op.inputs[0]
##return grad * 20 * x # add a "small" error just to see the difference:
def _MySquareGrad(op, grad):
#first get a Variable that was created using tf.get_variable()
with tf.variable_scope("", reuse=True):
var = tf.get_variable('var')
#now create the assign graph:
tmp = var*10.
assign_op = var.assign(tmp)
#now make the assign operation part of the grad calculation graph:
with tf.control_dependencies([assign_op]):
x = tf.identity(op.inputs[0])
return grad * 20 * x
with tf.Session() as sess:
x = tf.constant([1., 2.])
var = tf.get_variable(name="var", shape=[], initializer=tf.constant_initializer(0.2))
y = mysquare(x)
tf.global_variables_initializer().run()
print(x.eval(), y.eval(), tf.gradients(y, x)[0].eval())
print("Now var is 10 times larger:", var.eval())