SGFS clean up (#2427)

* [WIP] SGFS clean up cleaner init for SGFS * Update test * Remove unnecessary parameter
pymc-devs · Jul 19, 2017 · d06260d · d06260d
1 parent 99e3cd3
commit d06260d
Show file tree

Hide file tree

Showing 4 changed files with 134 additions and 249 deletions.
diff --git a/docs/source/notebooks/bayesian_neural_network_with_sgfs.ipynb b/docs/source/notebooks/bayesian_neural_network_with_sgfs.ipynb
diff --git a/docs/source/notebooks/sgfs_simple_optimization.ipynb b/docs/source/notebooks/sgfs_simple_optimization.ipynb
diff --git a/pymc3/step_methods/sgmcmc.py b/pymc3/step_methods/sgmcmc.py
@@ -98,18 +98,19 @@ def __init__(self,
                  batch_size=None,
                  total_size=None,
                  step_size=1.0,
-                 step_size_decay=100,
                  model=None,
                  random_seed=None,
                  minibatches=None,
                  minibatch_tensors=None,
                  **kwargs):
         warnings.warn(EXPERIMENTAL_WARNING)
-        if model is None:
-            model = modelcontext(model)
+
+        model = modelcontext(model)
 
         if vars is None:
-            vars = inputvars(model)
+            vars = model.vars
+
+        vars = inputvars(vars)
 
         self.model = model
         self.vars = vars
@@ -128,7 +129,6 @@ def __init__(self,
             self.random = tt_rng(random_seed)
 
         self.step_size = step_size
-        self.step_size_decay = step_size_decay
         shared = make_shared_replacements(vars, model)
         self.q_size = int(sum(v.dsize for v in self.vars))
 
@@ -237,7 +237,7 @@ def mk_training_fn(self):
         avg_I = self.avg_I
         t = self.t
         updates = self.updates
-        epsilon = self.step_size / pow(2.0, t // self.step_size_decay)
+        step_size = self.step_size
         random = self.random
         inarray = self.inarray
         gt, dlog_prior = self.dlogp_elemwise, self.dlog_prior
@@ -268,11 +268,11 @@ def mk_training_fn(self):
         # where B_ch is cholesky decomposition of B
         # i.e. B = dot(B_ch, B_ch^T)
         B_ch = tt.slinalg.cholesky(B)
-        noise_term = tt.dot((2.*B_ch)/tt.sqrt(epsilon), \
+        noise_term = tt.dot((2.*B_ch)/tt.sqrt(step_size), \
                 random.normal((q_size,), dtype=theano.config.floatX))
         # 9.
         # Inv. Fisher Cov. Matrix
-        cov_mat = (gamma * I_t * N) + ((4. / epsilon) * B)
+        cov_mat = (gamma * I_t * N) + ((4. / step_size) * B)
         inv_cov_mat = tt.nlinalg.matrix_inverse(cov_mat)
         # Noise Coefficient
         noise_coeff = (dlog_prior + (N * avg_gt) + noise_term)

diff --git a/pymc3/tests/test_sgfs.py b/pymc3/tests/test_sgfs.py
@@ -14,13 +14,14 @@ def f(x, a, b, c):
     a, b, c = 1, 2, 3
 
     batch_size = 50
-    x_train = np.random.uniform(-10, 10, size=(batch_size*500,)).astype('float32')
+    total_size = batch_size*500
+    x_train = np.random.uniform(-10, 10, size=(total_size,)).astype('float32')
     x_obs = pm.data.Minibatch(x_train, batch_size=batch_size)
 
     y_train = f(x_train, a, b, c) + np.random.normal(size=x_train.shape).astype('float32')
     y_obs = pm.data.Minibatch(y_train, batch_size=batch_size)
 
-    with Model() as model:
+    with Model():
         abc = Normal('abc', mu=mu0, sd=sd0, shape=(3,))
         x = x_obs
         x2 = x**2
@@ -29,7 +30,7 @@ def f(x, a, b, c):
         y = X.dot(abc)
         pm.Normal('y', mu=y, observed=y_obs)
 
-        step_method = pm.SGFS(vars=model.vars, batch_size=batch_size, step_size=1., total_size=draws*batch_size)
+        step_method = pm.SGFS(batch_size=batch_size, step_size=1., total_size=total_size)
         trace = pm.sample(draws=draws, step=step_method, init=None)
 
-    np.testing.assert_allclose(np.mean(trace['abc'], axis=0), np.asarray([a, b, c]), rtol=0.2)
+    np.testing.assert_allclose(np.mean(trace['abc'], axis=0), np.asarray([a, b, c]), rtol=0.1)