Модель Keras, вложенная в пользовательский слой Keras

Question

Я хотел бы построить модель Keras, которая использует числовой SPICE-подобный метод для прямого распространения.Поскольку проблемы SPICE аналитически не разрешимы, я создал следующий класс.Этот класс очень хорошо работает для реализации прогнозирования (численное прямое сокращение) и определения градиентов (аналитически).

Класс:

# "..." notes places where code is ommited for conciseness

class SPICE_solver():
    def __init__(self, num_inputs, num_outputs, ...):
    ...
        self.net = build_model_SPICE_solver(num_inputs, num_outputs, ...)

    def predict(self, activations, weights, ...):
        '''
        :param activations:         shape: (?, num_inputs)
        :param weights:             shape: (1, num_inputs, num_outputs)
        :return: vout               shape: (?, num_outputs)
        '''
        ...
        out = np.zeros([activations.shape[0], weights.shape[-1]])
        self.net.fit(x=[activations, weights],
                     y=[out],
                     epochs=200,
                     callbacks=[EarlyStoppingByLossVal(monitor='loss', value=self.acc, verbose=0)],
                     verbose=0,
                     steps_per_epoch=64)
        self.vout = self.net.get_weights()
        return self.vout    # weights incidate the output of the 'layer'

    def gradients(self, activations, weights, ...):
        '''
        :param activations:         shape: (?, num_inputs)
        :param weights:             shape: (?, num_inputs, num_outputs)
        :return: gradient:          list of gradients for: activations, weights (w.r.t. vout)
        '''
        ...
        outputTensor = self.net.output
        listOfVariableTensors = self.net.input
        gradients = K.gradients(outputTensor, listOfVariableTensors)
        sess = tf.InteractiveSession()
        sess.run(tf.initialize_all_variables())
        self.grad = sess.run(gradients, feed_dict={self.net.input[0]:activations, self.net.input[1]:weights})
        return self.grad

Я хотел бы использовать этот класс для выполнения прямого распространения (SPICE_solver.predict) и обратного распространения (SPICE_solver.gradients) в пользовательском слое Keras более высокого уровня.

Пользовательский слой Keras:

class mac_nonLin_SPICE(Layer):

    def __init__(self,
                 output_dim,
                 **kwargs):
        self.output_dim = output_dim
        super(mac_nonLin_SPICE, self).__init__(**kwargs)

    def build(self, input_shape):
        # Create a trainable weight variable for this layer.
        self.kernel = self.add_weight(name='kernel',
                                      shape=(1, int(input_shape[1]), self.output_dim),
                                      initializer='glorot_uniform',
                                      # constraint='UnitNorm',
                                      trainable=True)
        self.slvr = SPICE_solver(int(input_shape[1]), self.output_dim)
        super(mac_nonLin_SPICE, self).build(input_shape)  # Be sure to call this at the end

    def call(self, x):
        return self.slvr.predict(x, self.kernel)

    # def reutrn gradient():????
    # pass

    def compute_output_shape(self, input_shape):
        return (input_shape[0], self.output_dim)

У меня много проблем с вызовом моделей Keras во вложенном режиме.Есть ли практический способ реализовать такой объект в пользовательском слое Keras?

edit : моя интуиция подсказывает мне, что перестройка всего проекта с помощью низкоуровневых API TensorFlow является наиболее практичнойМетод, хотя и неудобный.Все еще в поисках простой работы Keras.

Любая помощь очень ценится!