我正在尝试创建一个自定义的Keras层,该层是用于在numpy数组上运行并返回numpy数组的函数的通用包装。我正在为Keras使用Tensorflow 2.0后端。
我认为这对于那些想要将某些外部代码实现为可训练的Keras层的人来说非常有用。
我的策略是将tf.numpy_fun
与@tf.custom_gradient
装饰器一起使用,以定义可包裹Tensorflow运算符的函数,该运算符包装了Numpy函数。原则上,如果有一种方法可以调用该函数,也可以有一个自定义函数来计算梯度,则应该可以通过反向传播对其进行训练。
在这里我展示了一个实现非线性变换的示例: y = B exp(-A x)我验证了张量流函数与使用{{ 1}}
但是当我尝试在自定义Keras层中实现该功能时,出现了一个我不完全理解的错误。这是我的代码:
tf.GradientTape()
运行此命令时,在import numpy as np
import tensorflow as tf
def myNumpyFun(B,A,x):
z = np.exp( - A.dot( x ) )
y = B.dot( z )
return (y,z)
def myNumpyFunGrad(B,A,x,z,dC_dyT):
tmp = dC_dyT
tmp = B.transpose().dot( tmp )
tmp = z*tmp
tmp = A.transpose().dot( tmp )
tmp = -tmp
dC_dxT = tmp
tmp = dC_dyT
tmp = B.transpose().dot( tmp )
tmp = z*tmp
XT = np.kron(np.eye(A.shape[0]),x)
tmp = XT.dot(tmp)
tmp = -tmp
dC_daT = tmp
dC_dAT = dC_daT.reshape(A.shape)
tmp = dC_dyT
ZT = np.kron(np.eye(B.shape[0]),z)
tmp = ZT.dot(tmp)
dC_dbT = tmp
dC_dBT = dC_dbT.reshape(B.shape)
return (dC_dBT, dC_dAT, dC_dxT)
@tf.custom_gradient
def myTensorflowFun(B,A,x):
(y,z) = tf.numpy_function(myNumpyFun, [B,A,x], tf.float32)
def grad(dC_dyT):
return tf.numpy_function(myNumpyFunGrad, [B,A,x,z,dC_dyT], tf.float32)
return y, grad
class MyKerasLayer(tf.keras.layers.Layer):
def __init__(self,shapeA,shapeB, **kwargs):
self.shapeA = shapeA
self.shapeB = shapeB
super(MyKerasLayer, self).__init__(**kwargs)
def build(self, input_shape):
self.A = self.add_weight(name='A',
shape=self.shapeA,
initializer='uniform',
trainable=True)
self.B = self.add_weight(name='B',
shape=self.shapeB,
initializer='uniform',
trainable=True)
super(MyKerasLayer, self).build(input_shape)
def call(self, x):
def fun(w):
return myTensorflowFun(self.B,self.A,w)
return tf.map_fn(fun, x, dtype=tf.float32)
def compute_output_shape(self, input_shape):
return self.call(tf.ones(input_shape)).shape
inputs = tf.keras.Input(shape=((2,1)),dtype=float)
outputs = MyKerasLayer([3,2],[4,3])(inputs)
model = tf.keras.Model(inputs=inputs, outputs=outputs)
opt = tf.keras.optimizers.Adam(learning_rate=0.01)
model.compile(loss='mean_squared_error', optimizer=opt)
N=10000
x_train = np.random.rand(N,2,1).astype(float)
x_test = np.random.rand(N,2,1).astype(float)
A = np.arange(6).astype(float).reshape([3,2])
B = np.arange(12).astype(float).reshape([4,3])
y_train = np.zeros([N,4,1],dtype=float)
y_test = np.zeros([N,4,1],dtype=float)
for iSamp in np.arange(N):
y_train[iSamp] = B.dot( np.exp(- A.dot( x_train[iSamp] ) ) )
y_test[iSamp] = B.dot( np.exp(- A.dot( x_test[iSamp] ) ) )
model.fit( x_train,
y_train,
batch_size=32,
epochs=5,
verbose=1,
validation_data=(x_test,y_test))
score = model.evaluate( x_test,
y_test,
verbose=0)
print('Test Loss: ', score)
weights = model.get_weights()
print('Trained Weights: ', weights)
内部的第一个tf.numpy_fun
调用中出现错误。这是错误消息:
OperatorNotAllowedInGraphError:在图形执行中不允许在MyTensorflowFun
上进行迭代。使用急切执行或使用@ tf.function装饰此功能。
我很困惑,因为我认为Tensorflow渴望在默认情况下启用执行。如果我运行tf.Tensor
,它将返回tf.executing_eagerly()
,但是如果我在错误发生之前运行它,它将返回True
。这是否意味着Keras在图形模式下执行?我该如何解决?
如果有人对在自定义Keras层中实现可训练的numpy函数的目标还有其他一般性建议,也很感兴趣。