Я пытаюсь определить модельную функцию для MCMC. Идея состоит в том, чтобы иметь смесь двух распределений, контролируемых с отношением вероятностей. Одна из моих попыток выглядела бы так:

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

root = tfd.JointDistributionCoroutine.Root

def model_fn():
    rv_p     = yield root(tfd.Sample(tfd.Uniform(0.0,1.0),1))

    catprobs = tf.stack([rv_p, 1.-rv_p],0)
    rv_cat = tfd.Categorical(probs=catprobs)

    rv_norm1  = tfd.Sample(tfd.Normal(0.0,1.0),1)
    rv_norm2  = tfd.Sample(tfd.Normal(3.0,1.0),1)

    rv_mix = yield tfd.Mixture(cat=rv_cat,
                     components=[
                        rv_norm1,
                        rv_norm2,
                     ])

jd = tfd.JointDistributionCoroutine(model_fn)
jd.sample(2)

Код не работает с:

ValueError: components[0] batch shape must be compatible with cat shape and other component batch shapes ((2, 2) vs ())

Не могли бы вы привести мне пример того, как использовать распределение Mixture таким образом, чтобы разрешить "любую" форму входов?

Я использую tensorflow 2.4.1 и tensorflow_probability 0.12.1 с python 3.6

0
rv_normal 27 Фев 2021 в 00:41

1 ответ

Лучший ответ

Я понял. Для справки вот пример кода:

import os
os.environ['CUDA_VISIBLE_DEVICES'] = '-1'
import tensorflow as tf
import tensorflow_probability as tfp
import matplotlib.pyplot as plt
tfd = tfp.distributions
tfb = tfp.bijectors

import numpy as np
from time import time

numdata = 10000
data = np.random.normal(0.0,1.0,numdata).astype(np.float32)
data[int(numdata/2):] = 0.0
_=plt.hist(data,30,density=True)

root = tfd.JointDistributionCoroutine.Root
def dist_fn(rv_p,rv_mu):
    rv_cat = tfd.Categorical(probs=tf.stack([rv_p, 1.-rv_p],-1))
    rv_norm  = tfd.Normal(rv_mu,1.0)
    rv_zero =  tfd.Deterministic(tf.zeros_like(rv_mu))
    
    rv_mix = tfd.Independent(
                tfd.Mixture(cat=rv_cat,
                            components=[rv_norm,rv_zero]),
                reinterpreted_batch_ndims=1)
    return rv_mix


def model_fn():
    rv_p    = yield root(tfd.Sample(tfd.Uniform(0.0,1.0),1))
    rv_mu   = yield root(tfd.Sample(tfd.Uniform(-1.,1. ),1))
    
    rv_mix  = yield dist_fn(rv_p,rv_mu)
    
jd = tfd.JointDistributionCoroutine(model_fn)
unnormalized_posterior_log_prob = lambda *args: jd.log_prob(args + (data,))

n_chains = 1

p_init = [0.3]
p_init = tf.cast(p_init,dtype=tf.float32)

mu_init = 0.1
mu_init = tf.stack([mu_init]*n_chains,axis=0)

initial_chain_state = [
    p_init,
    mu_init,
]

bijectors = [
    tfb.Sigmoid(),  # p
    tfb.Identity(),  # mu
]

step_size = 0.01

num_results = 50000
num_burnin_steps = 50000


kernel=tfp.mcmc.TransformedTransitionKernel(
    inner_kernel=tfp.mcmc.HamiltonianMonteCarlo(
    target_log_prob_fn=unnormalized_posterior_log_prob,
    num_leapfrog_steps=2,
    step_size=step_size,
    state_gradients_are_stopped=True),
    bijector=bijectors)

kernel = tfp.mcmc.SimpleStepSizeAdaptation(
    inner_kernel=kernel, num_adaptation_steps=int(num_burnin_steps * 0.8))

#XLA optim
@tf.function(autograph=False, experimental_compile=True)
def graph_sample_chain(*args, **kwargs):
  return tfp.mcmc.sample_chain(*args, **kwargs)


st = time()
trace,stats = graph_sample_chain(
      num_results=num_results,
      num_burnin_steps=num_burnin_steps,
      current_state=initial_chain_state,
      kernel=kernel)
et = time()
print(et-st)


ptrace, mutrace = trace
plt.subplot(121)
_=plt.hist(ptrace.numpy(),100,density=True)
plt.subplot(122)
_=plt.hist(mutrace.numpy(),100,density=True)
print(np.mean(ptrace),np.mean(mutrace))
0
rv_normal 28 Фев 2021 в 09:25