Tensorflow, py_func, or custom function

Question

I'm currently working on a quaternionic Neural Network using Tensorflow (I want to use GPUs). TensorFlow doesn't have support for quaternions, but you can represent than as a 4x4 real matrix, so it might be possible to build such a neural network in TensorFlow.

Is there a simple way to add a custom operation or to do a custom operation on tensors?

For example, I can write:

output_activation = tf.nn.softmax(tf.matmul(hidden_activation, Weight_to_ouput))

...and that's pretty cool! All you have to do is add a loss function and then do backpropagation. However, I want to do the same thing but with quaternions, for example:

output_activation = mySigmoid(myFunction(hidden_activation, Weight_to_output))

However, I need to transform the quaternions to and from tensors to optimize the GPU calculation. So I need to create a function that gets some tensors as parameters and returns the transformed tensors. I've looked at py_func, but it seems that you can't return tensors.

I tried the following, but it failed:

def layerActivation(inputTensor,WeightTensor):
    newTensor = tf.matmul(inputTensor,WeightTensor)
    return newTensor

...and in main():

x = placeholder ...
W_to_hidden = tf.Variable
test = tf.py_func(layerActivation, [x,_W_to_hidden], [tf.float32])

with tf.Session() as sess:
    tf.initialize_all_variables().run()
    king_return = sess.run(test, feed_dict={x: qtrain})

Error : Unimplemented: Unsupported object type Tensor

Ideally I could use this output_activation in the standard backprop algorithm of TensorFlow but I don't know if it's possible.

Answer 1

Depending on the functionality required, you might be able to implement your operation as a composition of existing TensorFlow ops, without needing to use tf.py_func().

For example, the following works and will run on a GPU:

def layer_activation(input_tensor, weight_tensor):
    return tf.matmul(input_tensor, weight_tensor)

# ...
x = tf.placeholder(...)
W_to_hidden = tf.Variable(...)
test = layer_activation(input_tensor, weight_tensor)
# ...

The main reason to use tf.py_func() is if your operations cannot be implemented using TensorFlow operations, and you want to inject some Python code (e.g. using NumPy) that works on the actual values of your tensor.

However, if your mySigmoid() or myFunction() operations cannot be implemented in terms of existing TensorFlow operations, and you want to implement them on GPU, then—as keveman says—you will need to add a new op.

Answer 2

If you want to run your custom operations on GPUs, you have to provide GPU implementation (kernels) in C++. Look at the documentation here for how to extend TensorFlow with custom operations, and especially the section on GPU support.