使用Tensorflow包装器实现Numpy功能的可训练自定义Keras图层

时间:2020-05-03 17:09:58

标签: numpy tensorflow keras

我正在尝试创建一个自定义的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函数的目标还有其他一般性建议,也很感兴趣。

0 个答案:

没有答案