Как вычислить все вторые производные (только диагональ матрицы Гессена) в тензорном потоке?

Question

У меня есть значение / функция потерь, и я хотел бы вычислить все вторые производные по тензору f (размера n). Мне удалось использовать tf.gradients дважды, но при повторном применении он суммирует производные по первому входу (см. Second_derivatives в моем коде).

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

import tensorflow as tf
import numpy as np

f = tf.Variable(np.array([[1., 2., 0]]).T)
loss = tf.reduce_prod(f ** 2 - 3 * f + 1)

first_derivatives = tf.gradients(loss, f)[0]

second_derivatives = tf.gradients(first_derivatives, f)[0]

hessian = [tf.gradients(first_derivatives[i,0], f)[0][:,0] for i in range(3)]

model = tf.initialize_all_variables()
with tf.Session() as sess:
    sess.run(model)
    print "\nloss\n", sess.run(loss)
    print "\nloss'\n", sess.run(first_derivatives)
    print "\nloss''\n", sess.run(second_derivatives)
    hessian_value = np.array(map(list, sess.run(hessian)))
    print "\nHessian\n", hessian_value

Я думал, что tf.gradients(first_derivatives, f[0, 0])[0] будет работать, например, для извлечения второй производной по f_0, но кажется, что тензорный поток не позволяет получить производную от тензора,

Answer 1

Теперь рассмотрим тф. Гессиан,

tf.hessians(loss, f)

https://www.tensorflow.org/api_docs/python/tf/hessians

Answer 2

tf.gradients([f1,f2,f3],...) вычисляет градиент f=f1+f2+f3Кроме того, дифференцируя по отношению к x[0] проблематично, потому что x[0] относится к новому Slice узел, который не является предком вашей потери, поэтому производная по нему будет None, Вы можете обойти это с помощью pack клеить x[0], x[1], ... вместе в xx и ваша потеря зависит от xx вместо x, Альтернативой является использование отдельных переменных для отдельных компонентов, в этом случае вычисление гессиана будет выглядеть примерно так.

def replace_none_with_zero(l):
  return [0 if i==None else i for i in l] 

tf.reset_default_graph()

x = tf.Variable(1.)
y = tf.Variable(1.)
loss = tf.square(x) + tf.square(y)
grads = tf.gradients([loss], [x, y])
hess0 = replace_none_with_zero(tf.gradients([grads[0]], [x, y]))
hess1 = replace_none_with_zero(tf.gradients([grads[1]], [x, y]))
hessian = tf.pack([tf.pack(hess0), tf.pack(hess1)])
sess = tf.InteractiveSession()
sess.run(tf.initialize_all_variables())
print hessian.eval()

Вот увидишь

[[ 2.  0.]
 [ 0.  2.]]

Answer 3

Следующая функция вычисляет 2-ю производную (диагональ матрицы Гессе) в Tensorflow 2.0:

%tensorflow_version 2.x  # Tells Colab to load TF 2.x
import tensorflow as tf

def calc_hessian_diag(f, x):
    """
    Calculates the diagonal entries of the Hessian of the function f
    (which maps rank-1 tensors to scalars) at coordinates x (rank-1
    tensors).
    
    Let k be the number of points in x, and n be the dimensionality of
    each point. For each point k, the function returns

      (d^2f/dx_1^2, d^2f/dx_2^2, ..., d^2f/dx_n^2) .

    Inputs:
      f (function): Takes a shape-(k,n) tensor and outputs a
          shape-(k,) tensor.
      x (tf.Tensor): The points at which to evaluate the Laplacian
          of f. Shape = (k,n).
    
    Outputs:
      A tensor containing the diagonal entries of the Hessian of f at
      points x. Shape = (k,n).
    """
    # Use the unstacking and re-stacking trick, which comes
    # from https://github.com/xuzhiqin1990/laplacian/
    with tf.GradientTape(persistent=True) as g1:
        # Turn x into a list of n tensors of shape (k,)
        x_unstacked = tf.unstack(x, axis=1)
        g1.watch(x_unstacked)

        with tf.GradientTape() as g2:
            # Re-stack x before passing it into f
            x_stacked = tf.stack(x_unstacked, axis=1) # shape = (k,n)
            g2.watch(x_stacked)
            f_x = f(x_stacked) # shape = (k,)
        
        # Calculate gradient of f with respect to x
        df_dx = g2.gradient(f_x, x_stacked) # shape = (k,n)
        # Turn df/dx into a list of n tensors of shape (k,)
        df_dx_unstacked = tf.unstack(df_dx, axis=1)

    # Calculate 2nd derivatives
    d2f_dx2 = []
    for df_dxi,xi in zip(df_dx_unstacked, x_unstacked):
        # Take 2nd derivative of each dimension separately:
        #   d/dx_i (df/dx_i)
        d2f_dx2.append(g1.gradient(df_dxi, xi))
    
    # Stack 2nd derivates
    d2f_dx2_stacked = tf.stack(d2f_dx2, axis=1) # shape = (k,n)
    
    return d2f_dx2_stacked

Вот пример использования с функцией f(x) = ln(r), где x являются трехмерными координатами и r - радиус сферических координат:

f = lambda q : tf.math.log(tf.math.reduce_sum(q**2, axis=1))
x = tf.random.uniform((5,3))

d2f_dx2 = calc_hessian_diag(f, x)
print(d2f_dx2)

Это будет выглядеть примерно так:

tf.Tensor(
[[ 1.415968    1.0215727  -0.25363517]
 [-0.67299247  2.4847088   0.70901346]
 [ 1.9416015  -1.1799507   1.3937857 ]
 [ 1.4748447   0.59702784 -0.52290654]
 [ 1.1786096   0.07442689  0.2396735 ]], shape=(5, 3), dtype=float32)

Мы можем проверить правильность реализации, вычислив лапласиан (т. Е. Суммируя диагональ матрицы Гессе) и сравнив с теоретическим ответом для выбранной нами функции, 2 / r^2:

print(tf.reduce_sum(d2f_dx2, axis=1)) # Laplacian from summing above results
print(2./tf.math.reduce_sum(x**2, axis=1)) # Analytic expression for Lapalcian

Получаю следующее:

tf.Tensor([2.1839054 2.5207298 2.1554365 1.5489659 1.49271  ], shape=(5,), dtype=float32)
tf.Tensor([2.1839058 2.5207298 2.1554365 1.5489662 1.4927098], shape=(5,), dtype=float32)

Они согласны с точностью до ошибки округления.

Answer 4

Если нет необходимости использовать Tensorflow для вычисления второй производной, вы можете сделать:

import sympy as sp
x =sp.symbols('x', real=True)
loss = x**2 - 3*x + 1
f = loss
a = sp.diff(f,x)
b = sp.diff(a)   
print ("First derivative of: loss = f ** 2 - 3 * f + 1 respect to f is: ", a,"and the 
second derivative is: ",b)

что приводит к:

First derivative of: loss = f ** 2 - 3 * f + 1 respect to f is:  2*x - 3 and the second derivative is:  2

Если вы хотите рассчитать гессиан, вы можете следовать этому примеру:

import sympy as sp
x,y,z = sp.symbols('x y z', real=True)
f = (x**2)*y*z

Jacobian = np.array([sp.diff(f,x),sp.diff(f,y),sp.diff(f,z)])
Hessian = 

np.matrix([[sp.diff(sp.diff(f,x),x),sp.diff(sp.diff(f,y),x),sp.diff(sp.diff(f,z),x)],

[sp.diff(sp.diff(f,x),y),sp.diff(sp.diff(f,y),y),sp.diff(sp.diff(f,z),y)],

[sp.diff(sp.diff(f,x),z),sp.diff(sp.diff(f,z),y),sp.diff(sp.diff(f,z),z)]]) 

print ("La Jacobian es:", Jacobian)
print ("La Hessian es:\n",Hessian)    

что приводит к:

La Jacobian es: [[2*x*y*z x**2*z x**2*y]] 
La Hessian es:
[[2*y*z 2*x*z 2*x*y]
 [2*x*z 0 x**2]
 [2*x*y x**2 0]]