Как рассчитать градиент дивергенции Кульбака-Лейблера двух распределений вероятности тензора относительно среднего значения распределения? - PullRequest
Новая
       38

Как рассчитать градиент дивергенции Кульбака-Лейблера двух распределений вероятности тензора относительно среднего значения распределения?

1 голос
/ 09 июля 2019

В tenorflow-2.0 я пытаюсь создать keras.layers.Layer, который выводит расхождение Кульбака-Лейблера (KL) между двумя tensorflow_probability.distributions. Я хотел бы рассчитать градиент выхода (то есть дивергенцию KL) по отношению к среднему значению одного из tensorflow_probability.distributions.

Во всех моих попытках, к сожалению, получаются градиенты 0.

Я попытался реализовать минимальный пример, показанный ниже. Мне было интересно, могут ли проблемы быть связаны с режимом нетерпеливого выполнения tf 2, поскольку я знаю аналогичный подход, который работал в tf 1, где нетерпеливое выполнение по умолчанию отключено.

Это минимальный пример, который я пробовал: 

import numpy as np
import tensorflow as tf
import tensorflow_probability as tfp
from tensorflow.keras.models import Model
from tensorflow.keras.layers import Layer,Input

# 1 Define Layer

class test_layer(Layer):

    def __init__(self, **kwargs):
        super(test_layer, self).__init__(**kwargs)

    def build(self, input_shape):
        self.mean_W = self.add_weight('mean_W',trainable=True)

        self.kernel_dist = tfp.distributions.MultivariateNormalDiag(
            loc=self.mean_W,
            scale_diag=(1.,)
        )
        super(test_layer, self).build(input_shape)

    def call(self,x):
        return tfp.distributions.kl_divergence(
            self.kernel_dist,
            tfp.distributions.MultivariateNormalDiag(
                loc=self.mean_W*0.,
                scale_diag=(1.,)
            )
        )

# 2 Create model

x = Input(shape=(3,))
fx = test_layer()(x)
test_model = Model(name='test_random', inputs=[x], outputs=[fx])


# 3 Calculate gradient

print('\n\n\nCalculating gradients: ')

# example data, only used as a dummy
x_data = np.random.rand(99,3).astype(np.float32)

for x_now in np.split(x_data,3):
#     print(x_now.shape)
    with tf.GradientTape() as tape:
        fx_now = test_model(x_now)
        grads = tape.gradient(
            fx_now,
            test_model.trainable_variables,
        )
        print('\nKL-Divergence: ', fx_now, '\nGradient: ',grads,'\n')

print(test_model.summary())

Вывод кода выше 

Calculating gradients: 

KL-Divergence:  tf.Tensor(0.0029436834, shape=(), dtype=float32) 
Gradient:  [<tf.Tensor: id=237, shape=(), dtype=float32, numpy=0.0>] 


KL-Divergence:  tf.Tensor(0.0029436834, shape=(), dtype=float32) 
Gradient:  [<tf.Tensor: id=358, shape=(), dtype=float32, numpy=0.0>] 


KL-Divergence:  tf.Tensor(0.0029436834, shape=(), dtype=float32) 
Gradient:  [<tf.Tensor: id=479, shape=(), dtype=float32, numpy=0.0>] 

Model: "test_random"
_________________________________________________________________
Layer (type)                 Output Shape              Param #   
=================================================================
input_1 (InputLayer)         [(None, 3)]               0         
_________________________________________________________________
test_layer_3 (test_layer)    ()                        1         
=================================================================
Total params: 1
Trainable params: 1
Non-trainable params: 0
_________________________________________________________________
None

Расхождение KL рассчитывается правильно, но результирующий градиент составляет 0. Какой будет правильный способ получения градиентов?

Ответы [ 2 ]

1 голос
/ 11 июля 2019

Мы работаем над распределениями и биекторами, делая их удобными для закрытия переменных в конструкторе. (Еще не сделали MVN.) Тем временем вы можете использовать tfd.Independent(tfd.Normal(loc=self.mean_W, scale=1), reinterpreted_batch_ndims=1), который, я думаю, будет работать в вашем методе сборки, потому что мы адаптировали Normal.

Также: вы видели пакет tfp.layers? В частности, https://www.tensorflow.org/probability/api_docs/python/tfp/layers/KLDivergenceAddLoss может быть вам интересно.

1 голос
/ 11 июля 2019

Если кому-то будет интересно, я узнал, как решить эту проблему:

Строка 

self.kernel_dist = tfp.distributions.MultivariateNormalDiag(
            loc=self.mean_W,
            scale_diag=(1.,)
        )

не должна быть внутри build() - метода определения класса слоя,а точнее внутри call() метода.Вот модифицированный пример: 

import numpy as np
import tensorflow as tf
import tensorflow_probability as tfp
from tensorflow.keras.models import Model
from tensorflow.keras.layers import Layer,Input

# 1 Define Layer

class test_layer(Layer):

    def __init__(self, **kwargs):
        super(test_layer, self).__init__(**kwargs)

    def build(self, input_shape):
        self.mean_W = self.add_weight('mean_W',trainable=True)
        super(test_layer, self).build(input_shape)

    def call(self,x):
        self.kernel_dist = tfp.distributions.MultivariateNormalDiag(
            loc=self.mean_W,
            scale_diag=(1.,)
        )
        return tfp.distributions.kl_divergence(
            self.kernel_dist,
            tfp.distributions.MultivariateNormalDiag(
                loc=self.mean_W*0.,
                scale_diag=(1.,)
            )
        )

# 2 Create model

x = Input(shape=(3,))
fx = test_layer()(x)
test_model = Model(name='test_random', inputs=[x], outputs=[fx])


# 3 Calculate gradient

print('\n\n\nCalculating gradients: ')

# example data, only used as a dummy
x_data = np.random.rand(99,3).astype(np.float32)

for x_now in np.split(x_data,3):
#     print(x_now.shape)
    with tf.GradientTape() as tape:
        fx_now = test_model(x_now)
        grads = tape.gradient(
            fx_now,
            test_model.trainable_variables,
        )
        print('\nKL-Divergence: ', fx_now, '\nGradient: ',grads,'\n')

print(test_model.summary())

Теперь вывод 



Calculating gradients: 

KL-Divergence:  tf.Tensor(0.024875917, shape=(), dtype=float32) 
Gradient:  [<tf.Tensor: id=742, shape=(), dtype=float32, numpy=0.22305119>] 


KL-Divergence:  tf.Tensor(0.024875917, shape=(), dtype=float32) 
Gradient:  [<tf.Tensor: id=901, shape=(), dtype=float32, numpy=0.22305119>] 


KL-Divergence:  tf.Tensor(0.024875917, shape=(), dtype=float32) 
Gradient:  [<tf.Tensor: id=1060, shape=(), dtype=float32, numpy=0.22305119>] 

Model: "test_random"
_________________________________________________________________
Layer (type)                 Output Shape              Param #   
=================================================================
input_2 (InputLayer)         [(None, 3)]               0         
_________________________________________________________________
test_layer_1 (test_layer)    ()                        1         
=================================================================
Total params: 1
Trainable params: 1
Non-trainable params: 0
_________________________________________________________________
None

, как и ожидалось.

Это то, что было изменено с tensorflow 1 на tensorflow 2

...