Спецификация полиномиальной модели в вероятности тензорного потока

Я играю со смешанной полиномиальной моделью дискретного выбора в Tensorflow Probability. Функция должна принимать входные данные из трех альтернатив. Выбранная альтернатива задается ВЫБРАННЫМ (тензор # наблюдений x3). Ниже приведено обновление кода, отражающее мой прогресс в решении проблемы (но проблема остается).

Сейчас я получаю сообщение об ошибке:

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

с обратной связью, предполагающей, что проблема заключается в вызове log_prob () для последнего компонента совместного распределения (т.е. tfp.Independent(tfp.Multinomial(...))

Основные компоненты (спасибо Падарну Уилсону за помощь в исправлении определения совместного распределения):

@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)

2 ответа

Решение

Я смог получить разумные результаты от своей модели. Всем спасибо за помощь! Следующие пункты помогли решить различные проблемы.

  1. Использование JointDistributionSequentialAutoBatched() для создания согласованных форм партии. Для доступа необходимо установить tf-nightly.

  2. Более информативные априоры для гиперпараметров. Экспоненциальное преобразование в распределении Multinomial() означает, что неинформативные гиперпараметры (т. Е. С sigma = 1e5) означают, что вы быстро вводите большие положительные числа в exp(), что приводит к бесконечности.

  3. Также важно было установить размер шага и т. Д.

  4. Я нашел ответ Кристофера Сутера на недавний вопрос на форуме Tensorflow Probability, в котором была указана полезная модель из STAN. Я взял выборку из своего предыдущего опыта в качестве отправной точки для полезных параметров начального правдоподобия.

  5. Несмотря на то, что JointDistributionSequentialAutoBatched() исправлял формы пакета, я вернулся и скорректировал свои совместные формы распределения, так что печать log_prob_parts() дает согласованные формы (например, [10,1] для 10 цепочек). Я все еще получаю ошибку формы без использования JointDistributionSequentialAutoBatched(), но комбинация, похоже, работает.

  6. Я разделил свой affine() на две функции. Они делают то же самое, но удаляют предупреждения об обратном отслеживании. По сути, affine() могла транслировать входные данные, но они различались, и было проще написать две функции, которые настраивают входные данные с согласованными формами. Входы разной формы заставляют Tensorflow отслеживать функцию несколько раз.

Скорее всего, это проблема с вашим исходным состоянием и количеством цепочек. Вы можете попробовать инициализировать ядро ​​вне вызова сэмплера:

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, и состояния предложения возвращаются.

Еще одна вещь, которую нужно сделать, - это передать ваше начальное состояние в ваш лог-вероятность / апостериор и посмотреть, соответствуют ли размеры возвращенных логарифмических вероятностей тому, что, по вашему мнению, должно быть (если вы выполняете несколько цепочек, тогда, возможно, он должен возвращать # цепочек логарифмическая вероятность).

Насколько я понимаю, размер партии (# цепочек) должен быть первым во всех ваших векторизованных вычислениях.

В самой последней части моего сообщения в блоге о тензорном потоке и настраиваемых вероятностях есть рабочий код для примера, который делает это.

Другие вопросы по тегам