Slow sampling of NegativeBinomial distribution

Aperta
  • Qualcosa non va?
#1,843 0 commenti 0 reazioni 0 assegnatari Vedi su GitHub

Nessuno ha ancora preso questa issue.

Valutazione

Difficoltà
4/5
Tempo stimato
3-5 giorni
Idoneità per principianti
35/100
Tipo di issue
Bug
Chiarezza
Abbastanza chiara
Stato di attività
Ferma
Stack tecnologico
python
Ambito
machine-learning

Direzione di ricerca

Inizia riproducendo il benchmark dell’issue con TensorFlow Probability 0.23.0, il JAX substrate e i tre valori di total_count. Traccia il punto di ingresso del sampling di NegativeBinomial e confronta il suo comportamento con l’esempio gamma-plus-Poisson di JAX; il lavoro è completato quando il caso con total_count piccolo presenta un miglioramento delle prestazioni verificato senza compromettere il comportamento di sampling riportato.

Scritto dal modello di indicizzazione a partire dal testo della issue.

Descrizione

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

Lingua principale
Jupyter Notebook
Stelle
4.4k
Fork
1.1k
Nessuna PR unita negli ultimi 30g

Preparare l'ambiente

Come iniziare

  1. Leggi tutta la issue e poi la guida ai contributi del progetto.
  2. Commenta sulla issue per dire che te ne occupi tu — evita che due persone facciano lo stesso lavoro.
  3. Fai un fork del repository e lavora su un branch.
  4. Apri una pull request che faccia riferimento al numero della issue.

Altre issue di tensorflow/probability

Tutte le issue di tensorflow/probability

Issue simili

Altre issue su Machine Learning

Ricevi le nuove issue nella tua casella

Un breve riepilogo di issue GitHub adatte ai principianti.