Distribuciones de TensorFlow: una suave introducción

Ver en TensorFlow.org Ejecutar en Google Colab Ver fuente en GitHub Descargar cuaderno

En este cuaderno, exploraremos las distribuciones de TensorFlow (TFD para abreviar). El objetivo de este portátil es que suba suavemente la curva de aprendizaje, incluida la comprensión del manejo de TFD de las formas de los tensores. Este cuaderno intenta presentar ejemplos antes que conceptos abstractos. Primero presentaremos formas canónicas fáciles de hacer las cosas y guardaremos la vista abstracta más general hasta el final. Si usted es el tipo que prefiere un tutorial más abstracto y de estilo de referencia, echa un vistazo a Entender TensorFlow Distribuciones formas . Si usted tiene alguna pregunta sobre el material aquí, no dude en contacto (o unirse a) la lista de distribución de probabilidad TensorFlow . Estamos felices de poder ayudar.

Antes de comenzar, necesitamos importar las bibliotecas adecuadas. Nuestra biblioteca general es tensorflow_probability . Por convención, generalmente nos referimos a la biblioteca distribuciones como tfd .

Tensorflow Eager es un entorno de ejecución imperativo para TensorFlow. En TensorFlow ansioso, cada operación de TF se evalúa inmediatamente y produce un resultado. Esto contrasta con el modo "gráfico" estándar de TensorFlow, en el que las operaciones TF agregan nodos a un gráfico que luego se ejecuta. Todo este cuaderno está escrito usando TF Eager, aunque ninguno de los conceptos presentados aquí se basa en eso, y TFP se puede usar en modo gráfico.

import collections

import tensorflow as tf
import tensorflow_probability as tfp
tfd = tfp.distributions

try:
  tf.compat.v1.enable_eager_execution()
except ValueError:
  pass

import matplotlib.pyplot as plt

Distribuciones univariadas básicas

Vamos a sumergirnos y crear una distribución normal:

n = tfd.Normal(loc=0., scale=1.)
n
<tfp.distributions.Normal 'Normal' batch_shape=[] event_shape=[] dtype=float32>

Podemos sacar una muestra de él:

n.sample()
<tf.Tensor: shape=(), dtype=float32, numpy=0.25322816>

Podemos sacar varias muestras:

n.sample(3)
<tf.Tensor: shape=(3,), dtype=float32, numpy=array([-1.4658079, -0.5653636,  0.9314412], dtype=float32)>

Podemos evaluar un problema de registro:

n.log_prob(0.)
<tf.Tensor: shape=(), dtype=float32, numpy=-0.9189385>

Podemos evaluar múltiples probabilidades logarítmicas:

n.log_prob([0., 2., 4.])
<tf.Tensor: shape=(3,), dtype=float32, numpy=array([-0.9189385, -2.9189386, -8.918939 ], dtype=float32)>

Disponemos de una amplia gama de distribuciones. Probemos un Bernoulli:

b = tfd.Bernoulli(probs=0.7)
b
<tfp.distributions.Bernoulli 'Bernoulli' batch_shape=[] event_shape=[] dtype=int32>
b.sample()
<tf.Tensor: shape=(), dtype=int32, numpy=1>
b.sample(8)
<tf.Tensor: shape=(8,), dtype=int32, numpy=array([1, 0, 0, 0, 1, 0, 1, 0], dtype=int32)>
b.log_prob(1)
<tf.Tensor: shape=(), dtype=float32, numpy=-0.35667497>
b.log_prob([1, 0, 1, 0])
<tf.Tensor: shape=(4,), dtype=float32, numpy=array([-0.35667497, -1.2039728 , -0.35667497, -1.2039728 ], dtype=float32)>

Distribuciones multivariadas

Crearemos una normal multivariante con una covarianza diagonal:

nd = tfd.MultivariateNormalDiag(loc=[0., 10.], scale_diag=[1., 4.])
nd
<tfp.distributions.MultivariateNormalDiag 'MultivariateNormalDiag' batch_shape=[] event_shape=[2] dtype=float32>

Comparando esto con la normal univariada que creamos anteriormente, ¿qué es diferente?

tfd.Normal(loc=0., scale=1.)
<tfp.distributions.Normal 'Normal' batch_shape=[] event_shape=[] dtype=float32>

Vemos que lo normal univariante tiene un event_shape de () , lo que indica que es una distribución de escalar. El multivariado normal tiene una event_shape de 2 , lo que indica la [espacio para eventos] básico (https://en.wikipedia.org/wiki/Event_ (probability_theory)) de esta distribución es de dos dimensiones.

El muestreo funciona igual que antes:

nd.sample()
<tf.Tensor: shape=(2,), dtype=float32, numpy=array([-1.2489667, 15.025171 ], dtype=float32)>
nd.sample(5)
<tf.Tensor: shape=(5, 2), dtype=float32, numpy=
array([[-1.5439653 ,  8.9968405 ],
       [-0.38730723, 12.448896  ],
       [-0.8697963 ,  9.330035  ],
       [-1.2541095 , 10.268944  ],
       [ 2.3475595 , 13.184147  ]], dtype=float32)>
nd.log_prob([0., 10])
<tf.Tensor: shape=(), dtype=float32, numpy=-3.2241714>

En general, las normales multivariadas no tienen covarianza diagonal. TFD ofrece múltiples formas de crear normales multivariadas, incluida una especificación de covarianza completa, que usamos aquí.

nd = tfd.MultivariateNormalFullCovariance(
    loc = [0., 5], covariance_matrix = [[1., .7], [.7, 1.]])
data = nd.sample(200)
plt.scatter(data[:, 0], data[:, 1], color='blue', alpha=0.4)
plt.axis([-5, 5, 0, 10])
plt.title("Data set")
plt.show()

png

Distribuciones múltiples

Nuestra primera distribución de Bernoulli representó un lanzamiento de una sola moneda justa. También podemos crear un lote de distribuciones de Bernoulli independientes, cada uno con sus propios parámetros, en una sola Distribution de objetos:

b3 = tfd.Bernoulli(probs=[.3, .5, .7])
b3
<tfp.distributions.Bernoulli 'Bernoulli' batch_shape=[3] event_shape=[] dtype=int32>

Es importante tener claro lo que esto significa. La llamada anterior define tres distribuciones de Bernoulli independientes, que pasan a ser contenida en el mismo Python Distribution objeto. Las tres distribuciones no se pueden manipular individualmente. Nota cómo el batch_shape es (3,) , lo que indica un lote de tres distribuciones, y la event_shape es () , que indica las distribuciones individuales tienen un espacio para eventos univariado.

Si llamamos a sample , se obtiene una muestra de los tres:

b3.sample()
<tf.Tensor: shape=(3,), dtype=int32, numpy=array([0, 1, 1], dtype=int32)>
b3.sample(6)
<tf.Tensor: shape=(6, 3), dtype=int32, numpy=
array([[1, 0, 1],
       [0, 1, 1],
       [0, 0, 1],
       [0, 0, 1],
       [0, 0, 1],
       [0, 1, 0]], dtype=int32)>

Si llamamos prob , (esto tiene la misma semántica de la forma como log_prob ; utilizamos prob con estos pequeños ejemplos de Bernoulli para mayor claridad, aunque log_prob normalmente se prefiere en aplicaciones) podemos pasar un vector y evaluar la probabilidad de cada moneda produciendo ese valor :

b3.prob([1, 1, 0])
<tf.Tensor: shape=(3,), dtype=float32, numpy=array([0.29999998, 0.5       , 0.29999998], dtype=float32)>

¿Por qué la API incluye forma de lote? Semánticamente, se podría realizar los mismos cálculos mediante la creación de una lista de distribución de e iterar sobre ellos con una for bucle (al menos en el modo de Eager, en modo gráfico TF se necesitaría una tf.while bucle). Sin embargo, tener un conjunto (potencialmente grande) de distribuciones idénticamente parametrizadas es extremadamente común, y el uso de cálculos vectorizados siempre que sea posible es un ingrediente clave para poder realizar cálculos rápidos utilizando aceleradores de hardware.

Uso de independientes para agregar lotes a eventos

En la sección anterior, hemos creado b3 , una sola Distribution objeto que representa tres lanzamientos de la moneda. Si llamamos b3.prob en un vector \(v\), la \(i\)'th entrada era la probabilidad de que la \(i\)ª moneda de toma valor \(v[i]\).

Supongamos que, en cambio, nos gustaría especificar una distribución "conjunta" sobre variables aleatorias independientes de la misma familia subyacente. Se trata de un objeto diferente matemáticamente, en que para esta nueva distribución, prob en un vector \(v\) devolverá un valor único que representa la probabilidad de que todo el conjunto de monedas coincide con el vector \(v\).

¿Cómo logramos esto? Utilizamos una distribución "de orden superior" llamada Independent , que tiene una distribución y produce una nueva distribución con la forma de lotes trasladado a la forma de evento:

b3_joint = tfd.Independent(b3, reinterpreted_batch_ndims=1)
b3_joint
<tfp.distributions.Independent 'IndependentBernoulli' batch_shape=[] event_shape=[3] dtype=int32>

Comparar la forma a la del original b3 :

b3
<tfp.distributions.Bernoulli 'Bernoulli' batch_shape=[3] event_shape=[] dtype=int32>

Como prometido, vemos que que Independent se ha movido la forma de proceso por lotes en la forma evento: b3_joint es una única distribución ( batch_shape = () ) durante un espacio para eventos tridimensional ( event_shape = (3,) ).

Revisemos la semántica:

b3_joint.prob([1, 1, 0])
<tf.Tensor: shape=(), dtype=float32, numpy=0.044999998>

Una forma alternativa de obtener el mismo resultado sería que las probabilidades de cómputo utilizando b3 y hacer la reducción de forma manual mediante la multiplicación (o, en el caso más habitual en el que se utilizan las probabilidades de registro, resumiendo):

tf.reduce_prod(b3.prob([1, 1, 0]))
<tf.Tensor: shape=(), dtype=float32, numpy=0.044999994>

Indpendent permite al usuario para representar de forma más explícita el concepto deseado. Consideramos que esto es extremadamente útil, aunque no es estrictamente necesario.

Hechos graciosos:

  • b3.sample y b3_joint.sample tienen diferentes implementaciones conceptuales, sino salidas indistinguibles: la diferencia entre un lote de distribuciones independientes y una única distribución creado a partir del lote usando Independent aparece cuando el cálculo de probabilités, no cuando el muestreo.
  • MultivariateNormalDiag podría implementarse trivialmente usando los escalares Normal y Independent distribuciones (no se aplica realmente de esta manera, pero podría ser).

Lotes de distorsiones multivariantes

Creemos un lote de tres normales multivariantes bidimensionales de covarianza completa:

nd_batch = tfd.MultivariateNormalFullCovariance(
    loc = [[0., 0.], [1., 1.], [2., 2.]],
    covariance_matrix = [[[1., .1], [.1, 1.]], 
                         [[1., .3], [.3, 1.]],
                         [[1., .5], [.5, 1.]]])
nd_batch
<tfp.distributions.MultivariateNormalFullCovariance 'MultivariateNormalFullCovariance' batch_shape=[3] event_shape=[2] dtype=float32>

Vemos batch_shape = (3,) , así que hay tres normales multivariados independientes, y event_shape = (2,) , por lo que cada normal multivariante es bidimensional. En este ejemplo, las distribuciones individuales no tienen elementos independientes.

Obras de muestreo:

nd_batch.sample(4)
<tf.Tensor: shape=(4, 3, 2), dtype=float32, numpy=
array([[[ 0.7367498 ,  2.730996  ],
        [-0.74080074, -0.36466932],
        [ 0.6516018 ,  0.9391426 ]],

       [[ 1.038303  ,  0.12231752],
        [-0.94788766, -1.204232  ],
        [ 4.059758  ,  3.035752  ]],

       [[ 0.56903946, -0.06875849],
        [-0.35127294,  0.5311631 ],
        [ 3.4635801 ,  4.565582  ]],

       [[-0.15989424, -0.25715637],
        [ 0.87479895,  0.97391707],
        [ 0.5211419 ,  2.32108   ]]], dtype=float32)>

Desde batch_shape = (3,) y event_shape = (2,) , se pasa un tensor de la forma (3, 2) a log_prob :

nd_batch.log_prob([[0., 0.], [1., 1.], [2., 2.]])
<tf.Tensor: shape=(3,), dtype=float32, numpy=array([-1.8328519, -1.7907217, -1.694036 ], dtype=float32)>

Radiodifusión, también conocida como ¿Por qué esto es tan confuso?

Haciendo abstracción de lo que hemos hecho hasta ahora, cada distribución tiene una forma de lotes B y un evento forma E . Vamos BE sea la concatenación de las formas de eventos:

  • Para las distribuciones escalares univariados n y b , BE = (). .
  • Para las normales multivariantes bidimensionales nd . BE = (2).
  • Por tanto b3 y b3_joint , BE = (3).
  • Para el lote de normales multivariantes ndb , BE = (3, 2).

Las "reglas de evaluación" que hemos estado usando hasta ahora son:

  • Muestra sin argumentos devuelve un tensor con forma de BE ; el muestreo con un escalar n devuelve un "n por BE " tensor.
  • prob y log_prob tomar un tensor de forma BE y devuelve un resultado de la forma B .

La "regla de evaluación" real para prob y log_prob es más complicada, de una manera que ofrece una potencia de potencial y velocidad, sino también la complejidad y los desafíos. La regla actual es (esencialmente) que el argumento a log_prob debe ser broadcastable contra BE ; cualquier dimensión "extra" se conserva en la salida.

Exploremos las implicaciones. Para el univariado normal de n , BE = () , por lo log_prob espera un escalar. Si pasamos log_prob un tensor con forma no vacía, los muestran como dimensiones de proceso por lotes en la salida:

n = tfd.Normal(loc=0., scale=1.)
n
<tfp.distributions.Normal 'Normal' batch_shape=[] event_shape=[] dtype=float32>
n.log_prob(0.)
<tf.Tensor: shape=(), dtype=float32, numpy=-0.9189385>
n.log_prob([0.])
<tf.Tensor: shape=(1,), dtype=float32, numpy=array([-0.9189385], dtype=float32)>
n.log_prob([[0., 1.], [-1., 2.]])
<tf.Tensor: shape=(2, 2), dtype=float32, numpy=
array([[-0.9189385, -1.4189385],
       [-1.4189385, -2.9189386]], dtype=float32)>

Turno de Let a la normal multivariante bidimensional nd (parámetros cambiados para fines ilustrativos):

nd = tfd.MultivariateNormalDiag(loc=[0., 1.], scale_diag=[1., 1.])
nd
<tfp.distributions.MultivariateNormalDiag 'MultivariateNormalDiag' batch_shape=[] event_shape=[2] dtype=float32>

log_prob "espera" una discusión con la forma (2,) , pero va a aceptar cualquier argumento de que las transmisiones en contra de esta forma:

nd.log_prob([0., 0.])
<tf.Tensor: shape=(), dtype=float32, numpy=-2.337877>

Pero podemos pasar "más" ejemplos, y evaluar toda su log_prob que está a la vez:

nd.log_prob([[0., 0.],
             [1., 1.],
             [2., 2.]])
<tf.Tensor: shape=(3,), dtype=float32, numpy=array([-2.337877 , -2.337877 , -4.3378773], dtype=float32)>

Quizás de manera menos atractiva, podemos transmitir sobre las dimensiones del evento:

nd.log_prob([0.])
<tf.Tensor: shape=(), dtype=float32, numpy=-2.337877>
nd.log_prob([[0.], [1.], [2.]])
<tf.Tensor: shape=(3,), dtype=float32, numpy=array([-2.337877 , -2.337877 , -4.3378773], dtype=float32)>

La transmisión de esta manera es una consecuencia de nuestro diseño de "habilitar la transmisión siempre que sea posible"; este uso es algo controvertido y podría eliminarse en una versión futura de TFP.

Ahora veamos el ejemplo de las tres monedas nuevamente:

b3 = tfd.Bernoulli(probs=[.3, .5, .7])

En este caso, el uso de la radiodifusión para representar la probabilidad de que cada moneda sale cara es bastante intuitiva:

b3.prob([1])
<tf.Tensor: shape=(3,), dtype=float32, numpy=array([0.29999998, 0.5       , 0.7       ], dtype=float32)>

(Compárese esto con b3.prob([1., 1., 1.]) , que tendríamos que volver usada en b3 se introdujo.)

Ahora supongamos que queremos saber, por cada moneda, la probabilidad de que la moneda sale cara y la probabilidad sale cruz. Podríamos imaginarnos intentando:

b3.log_prob([0, 1])

Desafortunadamente, esto produce un error con un seguimiento de pila largo y no muy legible. b3 tiene BE = (3) , por lo que debemos pasar b3.prob algo broadcastable contra (3,) . [0, 1] tiene la forma (2) , por lo que no difunde y crea un error. En cambio, tenemos que decir:

b3.prob([[0], [1]])
<tf.Tensor: shape=(2, 3), dtype=float32, numpy=
array([[0.7, 0.5, 0.3],
       [0.3, 0.5, 0.7]], dtype=float32)>

¿Por qué? [[0], [1]] tiene forma (2, 1) , por lo que difunde contra la forma (3) para hacer una forma de difusión de (2, 3) .

La transmisión es bastante poderosa: hay casos en los que permite una reducción de orden de magnitud en la cantidad de memoria utilizada y, a menudo, hace que el código de usuario sea más corto. Sin embargo, programar con. Si llama log_prob y obtener un error, una falta de difusión es casi siempre el problema.

Yendo más lejos

En este tutorial, hemos proporcionado (con suerte) una introducción sencilla. Algunos consejos para ir más allá:

  • event_shape , batch_shape y sample_shape pueden ser rango arbitrario (en este tutorial que siempre están bien escalar o rango 1). Esto aumenta la potencia, pero nuevamente puede conducir a desafíos de programación, especialmente cuando se trata de transmisión. Para una inmersión profunda adicional en manipulación de la forma, ver Comprensión TensorFlow Distribuciones formas .
  • TFP incluye un potente abstracción conocida como Bijectors , que en conjunción con TransformedDistribution , produce una forma flexible, de composición para crear fácilmente nuevas distribuciones que son transformaciones invertibles de las distribuciones existentes. Vamos a tratar de escribir un tutorial sobre esto pronto, pero mientras tanto, echa un vistazo a la documentación