How to use tf.function as a Class Method with Class Parameters, with Custom Gradient? - tensorflow2.0

I'm building a custom Keras layer that needs a custom gradient. This layer will internally hold a module that has the functionality I want. This module should have internal parameters that I can set from the outside, which affect the behavior of the tf function inside it.
This is the simplest way to trigger the issue:
import tensorflow as tf
class Module:
def __init__(self, arr_size):
self.arr_size = tf.convert_to_tensor(arr_size, dtype=tf.int32)
#tf.function
def grad(self):
return None
#tf.function
#tf.custom_gradient
def internal_function(self):
return self.arr_size, grad
def __call__(self, x):
return self.internal_function() + x
md = Module(4)
val = md(3)
I run into an error that says "Failed to convert elements of <Module> to Tensor. Consider casting elements to a supported type."
I take that to mean that tf.function can't take a self argument, where self is a class. But then, how can I parameterize the internals of a tf.function wrapped method? I previously used a builder method to return a tf.function wrapped method, but this ran into issues where it could not be saved using tf.keras.model.save_model because the wrapped tf.function had no name when added to the graph.

Related

Problems with saving custom model for TF serving

I defined simple custom model:
import tensorflow as tf
class CustomModule(tf.keras.layers.Layer):
def __init__(self):
super(CustomModule, self).__init__()
self.v = tf.Variable(1.)
def call(self, x):
print('Tracing with', x)
return x * self.v
def mutate(self, new_v):
self.v.assign(new_v)
I want to save it for serving and that is why I need to provide a function for “serving_default”. I’ve tried to do it like this:
module = CustomModule() module_with_signature_path = './tmp/1' call = tf.function(module.mutate, input_signature=[tf.TensorSpec([], tf.float32)]) tf.saved_model.save(module, module_with_signature_path, signatures=call)
I got an error:
ValueError: Got a non-Tensor value <tf.Operation 'StatefulPartitionedCall' type=StatefulPartitionedCall> for key 'output_0' in the output of the function __inference_mutate_8 used to generate the SavedModel signature 'serving_default'. Outputs for functions used as signatures must be a single Tensor, a sequence of Tensors, or a dictionary from string to Tensor.
How can I properly define signature while saving model? Thank you!

Using condition in tf2.1.0-keras

I am trying to using the bool condition in my own tf2.1.0-keras model, below is the simple example:
import tensorflow as tf
class TestKeras:
def __init__(self):
pass
def build_graph(self):
x = tf.keras.Input(shape=(2),batch_size=1)
x_value = x[0,0]
y = tf.cond(x_value > 0, lambda :tf.add(x_value,0), lambda :tf.add(x_value,0))
return tf.keras.models.Model(inputs=[x], outputs=[y])
if __name__ == "__main__":
tk = TestKeras()
model = tk.build_graph()
model.summary(line_length=100)
but it seem not work and throw the exception:
using a `tf.Tensor` as a Python `bool` is not allowed in Graph execution. Use Eager execution or decorate this function with #tf.function.
I have try to replace the tf.cond with tf.keras.backend.switch, but it still got the same error.
Also i have try to split the code y = tf.cond(xxx) into a single funtion and add the #tf.funcion decorator:
#tf.function
def compute_y(self,x):
return tf.cond(x > 0, lambda :tf.add(x,0), lambda :tf.add(x,0))
but it got another error:
Inputs to eager execution function cannot be Keras symbolic tensors, but found [<tf.Tensor 'strided_slice:0' shape=() dtype=float32>]
Anyone knows how can condition works in tf2.1.0-keras?
tf.keras.Input is a symbolic Tensor used to define an input for a keras model. Whenever you want to apply custom logic in a keras model, you should either subclass the Layer class, or use a Lambda layer.
For example, with a Lambda layer:
class TestKeras:
def __init__(self):
pass
def build_graph(self):
x = tf.keras.Input(shape=(2),batch_size=1)
def custom_fct(x):
x_value = x[0,0]
return tf.cond(x_value > 0, lambda :tf.add(x_value,0), lambda :tf.add(x_value,0))
y = tf.keras.layers.Lambda(custom_fct)(x)
return tf.keras.models.Model(inputs=[x], outputs=[y])

Keras Custom Batch Normalization layer with an extra variable that can be changed in run time

I have implemented a custom version of Batch Normalization with adding self.skip variable that act somehow as trainable. Here is the minimal code:
from tensorflow.keras.layers import BatchNormalization
import tensorflow as tf
# class CustomBN(tf.keras.layers.Layer):
class CustomBN(BatchNormalization):
def __init__(self, **kwargs):
super(CustomBN, self).__init__(**kwargs)
self.skip = False
def call(self, inputs, training=None):
if self.skip:
tf.print("I'm skipping")
else:
tf.print("I'm not skipping")
return super(CustomBN, self).call(inputs, training)
def build(self, input_shape):
super(CustomBN, self).build(input_shape)
To be crystal clear, all I have done so far are:
sub classing BatchNormalization: should I sub class tf.keras.layers.Layer?
defining self.skip to change the behavior of CustomBN layer in run time.
checking the state of self.skip in call method to act correspondingly.
Now, to change the behavior of the 'CustomBN' layer, I use
self.model.layers[ind].skip = state
where state is either True or False, and ind is the index number of CustomBN layer in the model.
the evident problem is that the value of self.skip will never change.
If you notice any mistakes please notify me.
By default, the call function in your layer will be called when the graph is built. Not on a per batch basis. Keras model compile method as a run_eagerly option that would cause your model to run (slower) in eager mode which would invoke your call function without building a graph. This is most likely not what you want to do however.
Ideally you want the flag that changes the behavior to be an input to the call method... For instance you can add an extra input to your graph which is simply this state flag and pass that to your layer.
The following is an example of how you can have a conditional graph on an extra parameter.
import numpy as np
import tensorflow as tf
from tensorflow import keras
from tensorflow.keras import layers
class MyLayerWithFlag(keras.layers.Layer):
def call(self, inputs, flag=None):
c_one = tf.constant([1], dtype=tf.float32)
if flag is not None:
x = tf.cond(
flag, lambda: tf.math.add(inputs, c_one),
lambda: inputs)
return x
return inputs
inputs = layers.Input(shape=(2,))
state = layers.Input(shape=(1,), dtype=tf.bool)
x = MyLayerWithFlag()(inputs, flag=state)
out = layers.Lambda(tf.reduce_sum)(x)
model = keras.Model([inputs, state], out)
data = np.array([[1., 2.]])
state = np.array([[True]])
model.predict((data, state))

Simple Custom Layer fails to build in simple model

When using a custom class to replace a lambda function, building a model fails.
I previously had this code in a lambda function and it worked fine but I was unable to save a model. I need to save the model that I'm building which has dependencies on this code snippet.
import keras
import tensorflow as tf
class ShapePositionLayer(keras.layers.Layer):
def call(self, x):
assert isinstance(x, list)
a, b = x
return keras.backend.gather(keras.backend.shape(a), b)
def compute_output_shape(self, input_shape):
return (1)
captions = keras.layers.Input(shape=[5,1024], name='captions')
batch_size = ShapePositionLayer()([captions,tf.constant(0,
dtype=tf.int32)])
model = keras.models.Model(inputs=[captions], outputs=[batch_size])
I expected to be able to build a model.
Instead receive error:
AttributeError: 'NoneType' object has no attribute '_inbound_nodes'

How to build this custom layer in Keras?

I'm building a NN that supports complex numbers. Currently working on complex activation. According to a Benjio paper, this is a good one:
Where b is a trainable parameter to be learnt. So I'm building a special layer to do this activation. I'm new to Keras and stuck already. I created this code below, but it gives an error with the build function. I have no idea what's happening, I just tried to copy the template. Please help.
class modrelu(Layer):
def __init__(self, **kwargs):
super(modrelu, self).__init__(**kwargs)
def build(self):
self.b= K.variable(value=np.random.rand()-0.5, dtype='float64')
super(modrelu, self).build() # Be sure to call this at the end
def call(self, x):
assert isinstance(x, list)
ip_r, ip_i = x
comp= tf.complex(ip_r, ip_i )
ABS= tf.math.abs(comp)
ANG= tf.math.angle(comp)
ABS= K.relu( self.b + ABS)
op_r= ABS * K.sin(angle) #K.dot ??
op_i= ABS * K.cos(angle)
return [op_r, op_i]
def compute_output_shape(self, input_shape):
assert isinstance(input_shape, list)
shape_a, shape_b = input_shape
return [shape_a, shape_b]
Comments on my code:
In the init I didn't add anything, cause it is an activation layer that takes no input when instantiated.
In the build method, I tried to add the b's. Not sure if I should use the self.add_weight method. Ideally, I want to have as many b's as the dimension of input.
In the call method, this one, I'm pretty sure what I'm doing. It is easy, I just implemented the function.
The last one, compute_output_shape, I just copied-pasted the template. The output should be the same as the input, cause it is just an activation layer.
Finally, the error for what its worth, I know it is nonsense
TypeError Traceback (most recent call last)
<ipython-input-5-3101a9226da5> in <module>
1 a=K.variable(np.array([1,2]))
2 b=K.variable(np.array([3,4]))
----> 3 act([a,b])
~\AppData\Local\conda\conda\envs\python36\lib\site-packages\keras\engine\base_layer.py in __call__(self, inputs, **kwargs)
429 'You can build it manually via: '
430 '`layer.build(batch_input_shape)`')
--> 431 self.build(unpack_singleton(input_shapes))
432 self.built = True
433
TypeError: build() takes 1 positional argument but 2 were given
There are several issues with your code.
First of all I should address the error you get from interpreter:
TypeError: build() takes 1 positional argument but 2 were given
The build method should take input_shape argument. Therefore you should declare build method as build(self, input_shape)
The second issue is undefined shape of the variables in the build method. You should explicitly declare shape of the variables. In your case the np.random.rand array should be of input_shape shape.
Another issue is that you are trying to return 2 results ([op_r, op_i]) in the call method. I'm not specialist in Keras but as far as I know you can't do it. Every Keras layer should have one and only one output. See here for the details: https://github.com/keras-team/keras/issues/3061
However if you use tensorflow backend you may use complex numbers (tf.complex) to return both real (op_r) and imagenary (op_i) parts of the complex number.
Here is the working implementation of modrelu layer with simple usage example. It is writtern for TensorFlow 1.12.0 which is distributed with it's own implementation of Keras API but I think you can easily adopt it for original Keras:
import tensorflow as tf
from tensorflow.python.keras import backend as K
from tensorflow.python.keras.engine import Layer
import numpy as np
class modrelu(Layer):
def __init__(self, **kwargs):
super(modrelu, self).__init__(**kwargs)
# provide input_shape argument in the build method
def build(self, input_shape):
# You should pass shape for your variable
self.b= K.variable(value=np.random.rand(*input_shape)-0.5,
dtype='float32')
super(modrelu, self).build(input_shape) # Be sure to call this at the end
def call(self, inputs, **kwargs):
assert inputs.dtype == tf.complex64
ip_r = tf.math.real(inputs)
ip_i = tf.math.imag(inputs)
comp = tf.complex(ip_r, ip_i )
ABS = tf.math.abs(comp)
ANG = tf.math.angle(comp)
ABS = K.relu(self.b + ABS)
op_r = ABS * K.sin(ANG) #K.dot ??
op_i = ABS * K.cos(ANG)
# return single tensor in the call method
return tf.complex(op_r, op_i)
real = tf.constant([2.25, 3.25])
imag = tf.constant([4.75, 5.75])
x = tf.complex(real, imag)
y = modrelu()(x)
with tf.Session() as sess:
sess.run(tf.global_variables_initializer())
print(sess.run(y))
P.S.: I didn't check the math so you should check it by yourself.
You are not coding the layer correctly, the build function takes a input_shape parameter, which you can use to initialize the weights/parameters of your layer.
You can see an example in Keras' source code.