Tensorflow概率多项式模型的规范

Question

我在处理Tensorflow概率中的混合多项式离散选择模型。该函数应该在三种备选方案中作出选择。选择的选项由选择的(一个# observationsx3张量)指定。下面是对代码的更新，以反映我在这个问题上的进展(但问题仍然存在)。

我目前得到了错误：

tensorflow.python.framework.errors_impl.InvalidArgumentError: Incompatible shapes: [6768,3] vs. [1,3,6768] [Op:Mul]

追溯提示问题在于调用log_prob()作为联合分发的最终组件(即tfp.Independent(tfp.Multinomial(.)。

主要组件是(感谢Padarn Wilson帮助修复联合发行定义)：

@tf.function
def affine(x, kernel_diag, bias=tf.zeros([])):
  """`kernel_diag * x + bias` with broadcasting."""
  kernel_diag = tf.ones_like(x) * kernel_diag
  bias = tf.ones_like(x) * bias
  return x * kernel_diag + bias

def mmnl_func():
    adj_AV_train = (tf.ones(num_idx) - AV[:,0]) * tf.constant(-9999.)
    adj_AV_SM = (tf.ones(num_idx) - AV[:,1]) * tf.constant(-9999.)
    adj_AV_car = (tf.ones(num_idx) - AV[:,2]) * tf.constant(-9999.)

    return tfd.JointDistributionSequential([
        tfd.Normal(loc=0., scale=1e5),  # mu_b_time
        tfd.HalfCauchy(loc=0., scale=5),  # sigma_b_time
        lambda sigma_b_time,mu_b_time: tfd.MultivariateNormalDiag(  # b_time
        loc=affine(tf.ones([num_idx]), mu_b_time[..., tf.newaxis]),
        scale_diag=sigma_b_time*tf.ones(num_idx)),
        tfd.Normal(loc=0., scale=1e5), # a_train
        tfd.Normal(loc=0., scale=1e5), # a_car
        tfd.Normal(loc=0., scale=1e5), # b_cost
        lambda b_cost,a_car,a_train,b_time: tfd.Independent(tfd.Multinomial(
          total_count=1,
          logits=tf.stack([
              affine(DATA[:,0], tf.gather(b_time, IDX[:,0], axis=-1), (a_train + b_cost * DATA[:,1] + adj_AV_train)),
              affine(DATA[:,2], tf.gather(b_time, IDX[:,0], axis=-1), (b_cost * DATA[:,3] + adj_AV_SM)),
              affine(DATA[:,4], tf.gather(b_time, IDX[:,0], axis=-1), (a_car + b_cost * DATA[:,5] + adj_AV_car))
          ], axis=1)
        ),reinterpreted_batch_ndims=1)
    ])

@tf.function
def mmnl_log_prob(mu_b_time, sigma_b_time, b_time, a_train, a_car, b_cost):
    return mmnl_func().log_prob(
      [mu_b_time, sigma_b_time, b_time, a_train, a_car, b_cost, CHOICE])

# Based on https://www.tensorflow.org/tutorials/customization/performance#python_or_tensor_args
# change constant values to tf.constant()
nuts_samples = tf.constant(1000)
nuts_burnin = tf.constant(500)
num_chains = tf.constant(1)
## Initial step size
init_step_size= tf.constant(0.3)
# Set the chain's start state.
initial_state = [
    tf.zeros([num_chains], dtype=tf.float32, name="init_mu_b_time"),
    tf.zeros([num_chains], dtype=tf.float32, name="init_sigma_b_time"),
    tf.zeros([num_chains, num_idx], dtype=tf.float32, name="init_b_time"),
    tf.zeros([num_chains], dtype=tf.float32, name="init_a_train"),
    tf.zeros([num_chains], dtype=tf.float32, name="init_a_car"),
    tf.zeros([num_chains], dtype=tf.float32, name="init_b_cost")
]

## NUTS (using inner step size averaging step)
##
@tf.function
def nuts_sampler(init):
    nuts_kernel = tfp.mcmc.NoUTurnSampler(
      target_log_prob_fn=mmnl_log_prob, 
      step_size=init_step_size,
      )
    adapt_nuts_kernel = tfp.mcmc.DualAveragingStepSizeAdaptation(
  inner_kernel=nuts_kernel,
  num_adaptation_steps=nuts_burnin,
  step_size_getter_fn=lambda pkr: pkr.step_size,
  log_accept_prob_getter_fn=lambda pkr: pkr.log_accept_ratio,
  step_size_setter_fn=lambda pkr, new_step_size: pkr._replace(step_size=new_step_size)
       )

    samples_nuts_, stats_nuts_ = tfp.mcmc.sample_chain(
  num_results=nuts_samples,
  current_state=initial_state,
  kernel=adapt_nuts_kernel,
  num_burnin_steps=tf.constant(100),
  parallel_iterations=tf.constant(5))
    return samples_nuts_, stats_nuts_

samples_nuts, stats_nuts = nuts_sampler(initial_state)

Answer 1

我从我的模型中得到了合理的结果。谢谢大家的帮助！以下几点有助于解决各种问题。

1. 使用JointDistributionSequentialAutoBatched()生成一致的批处理形状。你需要通宵安装才能进入。
2. 提供更多关于超参数的信息。多项式()分布中的指数变换意味着无信息的超参数(即σ= 1e5)意味着您很快就会有大量正数进入exp()，从而导致无穷大。
3. 设置步长等也很重要。
4. 我在Tensorflow概率论坛上找到了克里斯托弗·苏特( Christopher )最近提出的一个回答，该论坛指定了一个来自斯坦的模型。我使用了从我的先验抽取一个样本作为初始似然参数有用的起点。
5. 尽管JointDistributionSequentialAutoBatched()纠正了批处理形状，但我返回并更正了我的联合分发形状，以便打印log_prob_parts()提供一致的形状(即，10条链的10,1条)。在不使用JointDistributionSequentialAutoBatched()的情况下，我仍然会得到一个形状错误，但是这种组合似乎有效。
6. 我把我的仿射()分解成两个函数。他们做同样的事情，但删除了重跟踪警告。基本上，affine()能够广播输入，但它们是不同的，而且更容易编写两个函数来设置具有一致形状的输入。不同形状的输入导致Tensorflow多次跟踪函数。

Answer 2

更有可能的是，这是您的初始状态和链数的一个问题。您可以尝试在示例调用之外初始化内核：

nuts_kernel = tfp.mcmc.NoUTurnSampler(
      target_log_prob_fn=mmnl_log_prob, 
      step_size=init_step_size,
      )
    adapt_nuts_kernel = tfp.mcmc.DualAveragingStepSizeAdaptation(
  inner_kernel=nuts_kernel,
  num_adaptation_steps=nuts_burnin,
  step_size_getter_fn=lambda pkr: pkr.step_size,
  log_accept_prob_getter_fn=lambda pkr: pkr.log_accept_ratio,
  step_size_setter_fn=lambda pkr, new_step_size: pkr._replace(step_size=new_step_size)
       )

然后再做

nuts_kernel.bootstrap_results(initial_state)

并调查logL的形状，提案状态将被返回。

另一件要做的事情是将初始状态反馈到日志可能/后验中，并查看返回的日志概率的大小是否与您认为应该的大小相匹配(如果您正在执行多个链，那么它可能会返回# chains日志可能性)。

据我理解，批处理维度(#链)必须是所有矢量化计算中的第一个维度。

最后一部分我在tensorflow和定制可能性上的博客文章提供了一个这样做的示例的工作代码。