SVI now working with minibatches

GPy/core/svgp.py

@@ -5,9 +5,11 @@ import numpy as np
 from ..util import choleskies
 from sparse_gp import SparseGP
 from parameterization.param import Param
+from ..inference.latent_function_inference import SVGP as svgp_inf
+
 
 class SVGP(SparseGP):
-    def __init__(self, X, Y, Z, kernel, likelihood, name='SVGP', Y_metadata=None):
+    def __init__(self, X, Y, Z, kernel, likelihood, name='SVGP', Y_metadata=None, batchsize=None):
         """
         Stochastic Variational GP.
 
@@ -23,25 +25,33 @@ class SVGP(SparseGP):
 
         Hensman, Matthews and Ghahramani, Scalable Variational GP Classification, ArXiv 1411.2005
         """
+        if batchsize is None:
+            batchsize = X.shape[0]
+
+        self.X_all, self.Y_all = X, Y
+        # how to rescale the batch likelihood in case of minibatches
+        self.batchsize = batchsize
+        batch_scale = float(self.X_all.shape[0])/float(self.batchsize)
+        #KL_scale = 1./np.float64(self.mpi_comm.size)
+        KL_scale = 1.0
+
+        import climin.util
+        #Make a climin slicer to make drawing minibatches much quicker
+        self.slicer = climin.util.draw_mini_slices(self.X_all.shape[0], self.batchsize)
+        X_batch, Y_batch = self.new_batch()
 
         #create the SVI inference method
-        from ..inference.latent_function_inference import SVGP as svgp_inf
+        inf_method = svgp_inf(KL_scale=KL_scale, batch_scale=batch_scale)
-        inf_method = svgp_inf()
 
-        SparseGP.__init__(self,X, Y, Z, kernel, likelihood, inference_method=inf_method,
+        SparseGP.__init__(self, X_batch, Y_batch, Z, kernel, likelihood, inference_method=inf_method,
                  name=name, Y_metadata=Y_metadata, normalizer=False)
 
-        #?? self.set_data(X, Y)
-
         self.m = Param('q_u_mean', np.zeros((self.num_inducing, Y.shape[1])))
         chol = choleskies.triang_to_flat(np.tile(np.eye(self.num_inducing)[:,:,None], (1,1,Y.shape[1])))
         self.chol = Param('q_u_chol', chol)
         self.link_parameter(self.chol)
         self.link_parameter(self.m)
 
-        #self.batch_scale = 1. # how to rescale the batch likelihood in case of minibatches
-
-
     def parameters_changed(self):
         self.posterior, self._log_marginal_likelihood, self.grad_dict = self.inference_method.inference(self.q_u_mean, self.q_u_chol, self.kern, self.X, self.Z, self.likelihood, self.Y, self.Y_metadata)
 
@@ -55,16 +65,29 @@ class SVGP(SparseGP):
         if not self.Z.is_fixed:# only compute these expensive gradients if we need them
             self.Z.gradient = self.kern.gradients_X(self.grad_dict['dL_dKmm'], self.Z) + self.kern.gradients_X(self.grad_dict['dL_dKmn'], self.Z, self.X)
 
-
         self.likelihood.update_gradients(self.grad_dict['dL_dthetaL'])
         #update the variational parameter gradients:
         self.m.gradient = self.grad_dict['dL_dm']
         self.chol.gradient = self.grad_dict['dL_dchol']
 
+    def set_data(self, X, Y):
+        """
+        Set the data without calling parameters_changed to avoid wasted computation
+        If this is called by the stochastic_grad function this will immediately update the gradients
+        """
+        assert X.shape[1]==self.Z.shape[1]
+        self.X, self.Y = X, Y
 
-    #def set_data(self, X, Y):
+    def new_batch(self):
-        #assert X.shape[1]==self.Z.shape[1]
+        """
-        #self.X, self.Y = GPy.core.ObsAr(X), Y
+        Return a new batch of X and Y by taking a chunk of data from the complete X and Y
+        """
+        i = self.slicer.next()
+        return self.X_all[i], self.Y_all[i]
+
+    def stochastic_grad(self, parameters):
+        self.set_data(*self.new_batch())
+        return self._grads(parameters)
 
     def optimizeWithFreezingZ(self):
         self.Z.fix()
@@ -73,19 +96,3 @@ class SVGP(SparseGP):
         self.Z.unfix()
         self.kern.constrain_positive()
         self.optimize('bfgs')
-
-#class SPGPC_stoch(SPGPC):
-    #def __init__(self, X, Y, Z, kern=None, likelihood=None, batchsize=10):
-        #SPGPC.__init__(self, X[:1], Y[:1], Z, kern, likelihood)
-        #self.X_all, self.Y_all = X, Y
-        #self.batchsize = batchsize
-        #self.batch_scale = float(self.X_all.shape[0])/float(self.batchsize)
-#
-    #def stochastic_grad(self, w):
-        #i = np.random.permutation(self.X_all.shape[0])[:self.batchsize]
-        #self.set_data(self.X_all[i], self.Y_all[i])
-        #return self._grads(w)
-
-
-
-

GPy/inference/latent_function_inference/svgp.py

@@ -5,11 +5,14 @@ import numpy as np
 from posterior import Posterior
 
 class SVGP(LatentFunctionInference):
-    def inference(self, q_u_mean, q_u_chol, kern, X, Z, likelihood, Y, Y_metadata=None):
+    def __init__(self, KL_scale=1., batch_scale=1.):
-        assert Y.shape[1]==1, "multi outputs not implemented"
+        self.KL_scale = KL_scale
+        self.batch_scale = batch_scale
 
+    def inference(self, q_u_mean, q_u_chol, kern, X, Z, likelihood, Y, Y_metadata=None):
         num_inducing = Z.shape[0]
         num_data, num_outputs = Y.shape
+
         #expand cholesky representation
         L = choleskies.flat_to_triang(q_u_chol)
         S = np.einsum('ijk,ljk->ilk', L, L) #L.dot(L.T)
@@ -31,29 +34,25 @@ class SVGP(LatentFunctionInference):
         #compute the marginal means and variances of q(f)
         A = np.dot(Knm, Kmmi)
         mu = np.dot(A, q_u_mean)
-        #v = Knn_diag - np.sum(A*Knm,1) + np.sum(A*A.dot(S),1)
         v = Knn_diag[:,None] - np.sum(A*Knm,1)[:,None] + np.sum(A[:,:,None] * np.einsum('ij,jkl->ikl', A, S),1)
 
         #compute the KL term
         Kmmim = np.dot(Kmmi, q_u_mean)
-        #KL = -0.5*logdetS -0.5*num_inducing + 0.5*logdetKmm + 0.5*np.sum(Kmmi*S) + 0.5*q_u_mean.dot(Kmmim)
         KLs = -0.5*logdetS -0.5*num_inducing + 0.5*logdetKmm + 0.5*np.einsum('ij,ijk->k', Kmmi, S) + 0.5*np.sum(q_u_mean*Kmmim,0)
         KL = KLs.sum()
         dKL_dm = Kmmim
-        #dKL_dS = 0.5*(Kmmi - Si)
         dKL_dS = 0.5*(Kmmi[:,:,None] - Si)
-        #dKL_dKmm = 0.5*Kmmi - 0.5*Kmmi.dot(S).dot(Kmmi) - 0.5*Kmmim[:,None]*Kmmim[None,:]
         dKL_dKmm = 0.5*num_outputs*Kmmi - 0.5*Kmmi.dot(S.sum(-1)).dot(Kmmi) - 0.5*Kmmim.dot(Kmmim.T)
 
-        #if self.KL_scale:
+        KL_scale = self.KL_scale
-            #scale = 1./np.float64(self.mpi_comm.size)
+        batch_scale = self.batch_scale
-            #KL, dKL_dKmm, dKL_dS, dKL_dm = scale*KL, scale*dKL_dKmm, scale*dKL_dS, scale*dKL_dm
+        KL, dKL_dKmm, dKL_dS, dKL_dm = KL_scale*KL, KL_scale*dKL_dKmm, KL_scale*dKL_dS, KL_scale*dKL_dm
 
         #quadrature for the likelihood
         F, dF_dmu, dF_dv, dF_dthetaL = likelihood.variational_expectations(Y, mu, v)
 
         #rescale the F term if working on a batch
-        #F, dF_dmu, dF_dv =  F*batch_scale, dF_dmu*batch_scale, dF_dv*batch_scale
+        F, dF_dmu, dF_dv =  F*batch_scale, dF_dmu*batch_scale, dF_dv*batch_scale
 
         #derivatives of expected likelihood
         Adv = A.T[:,:,None]*dF_dv[None,:,:] # As if dF_Dv is diagonal
@@ -69,7 +68,6 @@ class SVGP(LatentFunctionInference):
         dF_dm = Admu
         dF_dS = AdvA
 
-
         #sum (gradients of) expected likelihood and KL part
         log_marginal = F.sum() - KL
         dL_dm, dL_dS, dL_dKmm, dL_dKmn = dF_dm - dKL_dm, dF_dS- dKL_dS, dF_dKmm- dKL_dKmm, dF_dKmn