Slow sampling of NegativeBinomial distribution

未关闭
  • 哪里不对？
#1,843 0 条评论 0 个 reaction 已指派 0 人 在 GitHub 查看

还没有人认领这个 Issue。

评估

难度
4/5
预计耗时
3-5 天
新手友好度
35/100
Issue 类型
缺陷
描述清晰度
基本清楚
活跃度
停滞
技术栈
python
领域
machine-learning

调研方向

首先，使用 TensorFlow Probability 0.23.0、JAX substrate 和三个 total_count 值复现 issue 中的基准测试。跟踪 NegativeBinomial 的采样入口点，并将其行为与 JAX 的 gamma-plus-Poisson 示例进行比较；当 total_count 较小的情况已验证性能得到改善，同时没有破坏报告中的采样行为时，即表示完成。

由索引模型根据 Issue 内容生成。

描述

Sampling from the Negative Binomial distribution (using jax substrates), especially using a small total_count parameter is very slow, compared to a jax only implementation.

import tensorflow_probability as tfp
tfp.__version__

'0.23.0'

from tensorflow_probability.substrates.jax.distributions import (
    NegativeBinomial as NBinom,
)
from jax import numpy as jnp, random as jrn, config as config
config.update("jax_enable_x64", True)

N = 1000
mu = 1e4
small_r = 0.1
middle_r = 10
large_r = 1000
key = jrn.PRNGKey(342354234)

nbinom_small_r = NBinom(total_count=small_r, logits=jnp.log(mu) - jnp.log(small_r))
nbinom_middle_r = NBinom(total_count=middle_r, logits=jnp.log(mu) - jnp.log(middle_r))
nbinom_large_r = NBinom(total_count=large_r, logits=jnp.log(mu) - jnp.log(large_r))

%timeit nbinom_small_r.sample(seed=key, sample_shape=(N,)).block_until_ready()
%timeit nbinom_middle_r.sample(seed=key, sample_shape=(N,)).block_until_ready()
%timeit nbinom_large_r.sample(seed=key, sample_shape=(N,)).block_until_ready()

6 s ± 38 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)
235 ms ± 710 µs per loop (mean ± std. dev. of 7 runs, 1 loop each)
119 ms ± 87.5 µs per loop (mean ± std. dev. of 7 runs, 10 loops each)

def sample_nbinom(key, r, mu, shp):
    key, sk_g, sk_p = jrn.split(key, 3)
    gamma_sample = mu / r * jrn.gamma(sk_g, r, shp)
    return jrn.poisson(sk_p, gamma_sample)
%timeit sample_nbinom(jrn.PRNGKey(0), small_r, mu, (N,)).block_until_ready()
%timeit sample_nbinom(jrn.PRNGKey(0), middle_r, mu, (N,)).block_until_ready()
%timeit sample_nbinom(jrn.PRNGKey(0), large_r, mu, (N,)).block_until_ready()

747 µs ± 65.8 µs per loop (mean ± std. dev. of 7 runs, 1 loop each)
604 µs ± 38 µs per loop (mean ± std. dev. of 7 runs, 1 loop each)
532 µs ± 867 ns per loop (mean ± std. dev. of 7 runs, 1,000 loops each)

I'm using tfp 0.23.0 under Python 3.10.13, as sampling under 0.24.0 with Python 3.12. does not work for me (I encounter similar behavior as in #1838).

主要语言
Jupyter Notebook
星标
4.4k
派生
1.1k
30 天内没有已合并 PR

环境准备

  • 没有 Dockerfile 或 Docker Compose 文件
  • 没有 Pull Request 模板
  • 阅读贡献指南

从这里开始

  1. 先读完整个 Issue，再读项目的贡献指南。
  2. 在 Issue 下留言说明你要接手 —— 这能避免两个人做同样的事。
  3. Fork 仓库，在一个分支上完成修改。
  4. 提交 Pull Request，并在描述里引用这个 Issue 编号。

tensorflow/probability 的其他 Issue

查看 tensorflow/probability 的全部 Issue

相似的 Issue

更多 Machine Learning Issue

把新 issue 发到你的邮箱

精选适合新手参与的 GitHub issue 摘要。