Проблема вычисления частных производных с GradientTape () в TensorFlow2

Question

У меня проблемы с вычислением градиентов с помощью автоматического дифференцирования c в TensorFlow. По сути, я хочу создать нейронную сеть, которая имеет только одно выходное значение f и получить входные данные двух значений (x, t). Сеть должна действовать как математическая функция, поэтому в этом случае f (x, t), где x и t - входные переменные, и я хочу вычислить частные производные, например, df_dx, d2f/dx2 или df_dt. Эти частные производные мне понадобятся позже для конкретной c функции потерь. Вот мой упрощенный код:

import numpy as np
import tensorflow as tf 
from tensorflow.keras.layers import Dense, Flatten
from tensorflow.keras import Model


class MyModel(Model):
    def __init__(self):
        super(MyModel, self).__init__()
        self.flatten = Flatten(input_shape=(2, 1))
        self.d1 = Dense(28)
        self.f = Dense(1)

    def call(self, y):
        y = self.flatten(y)
        y = self.d1(y)
        y = self.f(y)
        return y

if __name__ == "__main__":

    #inp contains the input-variables (x,t)
    inp = np.random.rand(1,2,1)
    inp_tf = tf.convert_to_tensor(inp, np.float32)   

    #Create a Model
    model = MyModel()

    #Here comes the important part:
    x = inp_tf[0][0]
    t = inp_tf[0][1]

    with tf.GradientTape(persistent=True) as tape:
        tape.watch(inp_tf[0][0])
        tape.watch(inp_tf)
        f = model(inp_tf)

    df_dx = tape.gradient(f, inp_tf[0][0])  #Derivative df_dx
    grad_f = tape.gradient(f, inp_tf)

    tf.print(f)         #--> [[-0.0968768075]]
    tf.print(df_dx)     #--> None
    tf.print(grad_f)    #--> [[[0.284864038]
                        #      [-0.243642956]]]

Я ожидал, что получу df_dx = [0.284864038] (первый компонент grad_f), но это приведет к None. Мои вопросы:

1. Можно ли получить частичные производные от f только для одной входной переменной?
2. Если да: что мне нужно изменить в моем коде, чтобы вычисление df_dx не приводит None?

Что я мог бы сделать, это изменить архитектуру class MyModel, в которой я использую два разных Inputlayer (один для x и один для t) так, чтобы я Я могу назвать модель как f = model(x,t), но это кажется мне неестественным, и я думаю, что должен быть более простой способ.

---

Другой момент заключается в том, что я не получаю сообщение об ошибке при изменении input_shape Flattenlayer, например, к self.flatten = Flatten(input_shape=(5,1), но мой inputvector имеет форму (1,2,1), поэтому я ожидаю получить ошибку, но это не так, почему? Я благодарен за вашу помощь:)

---

Я использую следующие конфигурации:

• Код Visual Studio с Python - расширение в качестве IDE
• Python -Версия: 3.7.6
• TensorFlow-версия: 2.1.0
• Keras-версия: 2.2.4-tf

Answer 1

Каждый раз, когда вы делаете inp_tf[0][0] или inp_tf[0][1], вы создаете новый тензор, но этот новый тензор не используется в качестве входных данных для вашей модели, inp_tf - это. Даже если inp_tf[0][0], если является частью inp_tf, с точки зрения TensorFlow нет графа вычислений между вновь созданными inp_tf[0][0] и f, следовательно, градиента нет. Вы должны вычислить градиент по отношению к inp_tf, а затем взять нужные вам части градиента.

В дополнение к этому, как показано в документации tf.GradientTape, вы можете использовать вложенные ленты для вычисления производных второго порядка. И, если вы используете jacobian, вы можете избежать использования persistent=True, что лучше для производительности. Вот как это может работать в вашем примере (я изменил функции активации слоя на sigmoid, поскольку линейная активация по умолчанию не будет иметь производную второго порядка).

import numpy as np
import tensorflow as tf 
from tensorflow.keras.layers import Dense, Flatten
from tensorflow.keras import Model

class MyModel(Model):
    def __init__(self):
        super(MyModel, self).__init__()
        self.flatten = Flatten(input_shape=(2, 1))
        self.d1 = Dense(28, activation='sigmoid')
        self.f = Dense(1, activation='sigmoid')

    def call(self, y):
        y = self.flatten(y)
        y = self.d1(y)
        y = self.f(y)
        return y

np.random.seed(0)
inp = np.random.rand(1, 2, 1)
inp_tf = tf.convert_to_tensor(inp, np.float32)
model = MyModel()
with tf.GradientTape() as tape:
    tape.watch(inp_tf)
    with tf.GradientTape() as tape2:
        tape2.watch(inp_tf)
        f = model(inp_tf)
    grad_f = tape2.gradient(f, inp_tf)
    df_dx = grad_f[0, 0]
    df_dt = grad_f[0, 1]
j = tape.jacobian(grad_f, inp_tf)
d2f_dx2 = j[0, 0, :, 0, 0]
d2f_dyx = j[0, 0, :, 0, 1]
d2f_dy2 = j[0, 1, :, 0, 1]
d2f_dxy = j[0, 1, :, 0, 0]

tf.print(df_dx)
# [0.0104712956]
tf.print(df_dt)
# [-0.00301733566]
tf.print(d2f_dx2)
# [[-0.000243180315]]
tf.print(d2f_dyx)
# [[-0.000740956515]]
tf.print(d2f_dy2)
# [[1.49392872e-05]]
tf.print(d2f_dxy)
# [[-0.000740956573]]