einsumming in rbf for speed

This commit is contained in:
James Hensman 2014-02-28 14:20:17 +00:00
parent d5658660a6
commit c87bda9e49

View file

@ -70,6 +70,8 @@ class RBF(Stationary):
self.lengthscale.gradient += (dL_dpsi2[:,:,:,None] * _dpsi2_dlengthscale).reshape(-1,self.input_dim).sum(axis=0) self.lengthscale.gradient += (dL_dpsi2[:,:,:,None] * _dpsi2_dlengthscale).reshape(-1,self.input_dim).sum(axis=0)
return return
elif isinstance(variational_posterior, variational.NormalPosterior):
l2 = self.lengthscale **2 l2 = self.lengthscale **2
#contributions from psi0: #contributions from psi0:
@ -99,6 +101,9 @@ class RBF(Stationary):
self.variance.gradient += 2.*np.sum(dL_dpsi2 * psi2)/self.variance self.variance.gradient += 2.*np.sum(dL_dpsi2 * psi2)/self.variance
else:
raise ValueError, "unknown distriubtion received for psi-statistics"
def gradients_Z_expectations(self, dL_dpsi1, dL_dpsi2, Z, variational_posterior): def gradients_Z_expectations(self, dL_dpsi1, dL_dpsi2, Z, variational_posterior):
# Spike-and-Slab GPLVM # Spike-and-Slab GPLVM
if isinstance(variational_posterior, variational.SpikeAndSlabPosterior): if isinstance(variational_posterior, variational.SpikeAndSlabPosterior):
@ -113,13 +118,13 @@ class RBF(Stationary):
return grad return grad
elif isinstance(variational_posterior, variational.NormalPosterior):
l2 = self.lengthscale **2 l2 = self.lengthscale **2
#psi1 #psi1
denom, dist, dist_sq, psi1 = self._psi1computations(Z, variational_posterior) denom, dist, dist_sq, psi1 = self._psi1computations(Z, variational_posterior)
denominator = l2 * denom grad = np.einsum('ij,ij,ijk,ijk->jk', dL_dpsi1, psi1, dist, -1./(denom*l2))
dpsi1_dZ = -psi1[:, :, None] * (dist / denominator)
grad = np.sum(dL_dpsi1[:, :, None] * dpsi1_dZ, 0)
#psi2 #psi2
Zdist, Zdist_sq, mudist, mudist_sq, psi2 = self._psi2computations(Z, variational_posterior) Zdist, Zdist_sq, mudist, mudist_sq, psi2 = self._psi2computations(Z, variational_posterior)
@ -127,10 +132,11 @@ class RBF(Stationary):
S = variational_posterior.variance S = variational_posterior.variance
term2 = mudist / (2.*S[:,None,None,:] + l2) # N, M, M, Q term2 = mudist / (2.*S[:,None,None,:] + l2) # N, M, M, Q
dZ = psi2[:, :, :, None] * (term1[None, :, :, :] + term2) #N,M,M,Q grad += 2.*np.einsum('ijk,ijk,ijkl->kl', dL_dpsi2, psi2, term1[None,:,:,:] + term2)
grad += 2*(dL_dpsi2[:, :, :, None] * dZ).sum(0).sum(0)
return grad return grad
else:
raise ValueError, "unknown distriubtion received for psi-statistics"
def gradients_qX_expectations(self, dL_dpsi0, dL_dpsi1, dL_dpsi2, Z, variational_posterior): def gradients_qX_expectations(self, dL_dpsi0, dL_dpsi1, dL_dpsi2, Z, variational_posterior):
# Spike-and-Slab GPLVM # Spike-and-Slab GPLVM
@ -151,6 +157,8 @@ class RBF(Stationary):
return grad_mu, grad_S, grad_gamma return grad_mu, grad_S, grad_gamma
elif isinstance(variational_posterior, variational.NormalPosterior):
l2 = self.lengthscale **2 l2 = self.lengthscale **2
#psi1 #psi1
denom, dist, dist_sq, psi1 = self._psi1computations(Z, variational_posterior) denom, dist, dist_sq, psi1 = self._psi1computations(Z, variational_posterior)
@ -161,8 +169,11 @@ class RBF(Stationary):
_, _, mudist, mudist_sq, psi2 = self._psi2computations(Z, variational_posterior) _, _, mudist, mudist_sq, psi2 = self._psi2computations(Z, variational_posterior)
S = variational_posterior.variance S = variational_posterior.variance
tmp = psi2[:, :, :, None] / (2.*S[:,None,None,:] + l2) tmp = psi2[:, :, :, None] / (2.*S[:,None,None,:] + l2)
grad_mu += -2.*(dL_dpsi2[:, :, :, None] * tmp * mudist).sum(1).sum(1) grad_mu += -2.*np.einsum('ijk,ijkl,ijkl->il', dL_dpsi2, tmp , mudist)
grad_S += (dL_dpsi2[:, :, :, None] * tmp * (2.*mudist_sq - 1)).sum(1).sum(1) grad_S += np.einsum('ijk,ijkl,ijkl->il', dL_dpsi2 , tmp , (2.*mudist_sq - 1))
else:
raise ValueError, "unknown distriubtion received for psi-statistics"
return grad_mu, grad_S return grad_mu, grad_S