在Keras中使用自定义Tensorflow Ops

时间:2018-08-01 16:20:21

标签: python tensorflow keras

我在tensorflow中有一个脚本,其中包含自定义tensorflow操作。我想将代码移植到keras,但不确定如何在keras代码中调用自定义操作。

我想在keras中使用tensorflow,所以到目前为止我发现的教程描述了与我想要的相反的东西:https://blog.keras.io/keras-as-a-simplified-interface-to-tensorflow-tutorial.html

我还阅读了有关可以包装任意自定义函数的Lambda层的信息,但是我没有看到tf.ops的示例。

如果您能提供一个带有最简单示例的代码片段,我将不胜感激。例如,假设tf.ops为:

outC = my_custom_op(inA, inB)

---编辑: here中描述了类似的问题-本质上是在keras中调用this自定义op,但是我无法掌握如何将其应用于我想要的另一个示例(例如this)上的解决方案。这个自定义tf op首先被编译(针对gpu),然后到目前为止在here中在tensorflow中使用,请参阅@第40行。对我来说,很清楚如何使用包装在Lambda层中的自定义(lambda)函数,我想了解的是,如果使用keras,如何使用编译后的自定义操作。

1 个答案:

答案 0 :(得分:4)

您可以将任意张量流函数包装在keras Lambda层中并将其添加到模型中。最小的工作示例from this answer

import tensorflow as tf
from keras.layers import Dense, Lambda, Input
from keras.models import Model

W = tf.random_normal(shape=(128,20))
b = tf.random_normal(shape=(20,))

inp = Input(shape=(10,))
x = Dense(128)(inp)
# Custom linear transformation
y = Lambda(lambda x: tf.matmul(x, W) + b, name='custom_layer')(x) 
model = Model(inp, y)