toy_minibatch_100_000.py

#!/usr/bin/env python
"""
This is just to show how fast we can do a minibatch gradient descent
of 100,000 samples (!).
"""
import edward as ed
import tensorflow as tf

from edward.stats import bernoulli
from edward.variationals import Variational, Bernoulli

class BernoulliModel:
    """
    p(x, z) = p(z) = p(z | x) = Bernoulli(z; p)
    """
    def __init__(self, p):
        self.p = p

    def log_prob(self, xs, zs):
        return bernoulli.logpmf(zs, p)

ed.set_seed(42)
p = tf.constant(0.6)
model = BernoulliModel(p)
variational = Variational()
variational.add(Bernoulli(1))

inference = ed.MFVI(model, variational)
inference.run(n_minibatch=int(1e5))