Обучаемый пользовательский слой Keras с использованием оболочки Tensorflow для функций Numpy

Question

Я пытаюсь создать пользовательский слой Keras, который является оберткой общего назначения для функций, которые работают с массивами numpy и возвращают массивы numpy. Я использую серверную часть Tensorflow 2.0 для Keras.

Я думаю, что это может быть очень полезно для тех, у кого есть какой-то внешний код, который они хотят реализовать в качестве обучаемого слоя Keras.

Моя стратегия заключалась в том, чтобы использовать tf.numpy_fun вместе с декоратором @tf.custom_gradient, чтобы определить дифференцируемый оператор Tensorflow, который заключает в себе функцию Numpy. В принципе, если есть способ вызова функции, а также пользовательская функция для вычисления градиентов, она должна быть обучаемой с помощью обратного распространения.

Здесь я показываю пример реализации нелинейного преобразования: y = B exp (-A x) Я проверил, что функция тензорного потока совместима с автоматическим дифференцированием Tensorflow c, используя tf.GradientTape()

Но когда я пытаюсь реализовать функцию в пользовательском Слой Keras Я получаю ошибку, которую не до конца понимаю. Вот мой код:

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 вызове внутри MyTensorflowFun. Вот сообщение об ошибке:

OperatorNotAllowedInGraphError: итерации по tf.Tensor недопустимы при выполнении Graph. Используйте Eager выполнение или украсьте эту функцию с помощью @ tf.function.

Я запутался, потому что я думал, что Tensorflow имеет активное выполнение по умолчанию. Если я запускаю tf.executing_eagerly(), он возвращает True, но если я запускаю его прямо перед ошибкой, он возвращает False. Значит ли это, что Keras выполняется в графическом режиме? Как я могу это исправить?

Также интересно, есть ли у кого-то другие общие предложения для достижения цели обучаемых numpy функций в пользовательском слое Keras.