class TestProblemBackendInterface(TestCase):
def setUp(self):
self.m = m = 20
self.n = n = 10
self.rank = rank = 3
A = np.random.randn(m, n)
@pymanopt.function.Autograd
def cost(u, s, vt, x):
return np.linalg.norm(((u * s) @ vt - A) @ x)**2
self.cost = cost
self.gradient = self.cost.compute_gradient()
self.hvp = self.cost.compute_hessian_vector_product()
self.manifold = Product([FixedRankEmbedded(m, n, rank), Euclidean(n)])
self.problem = pymanopt.Problem(self.manifold, self.cost)
def test_cost_function(self):
(u, s, vt), x = self.manifold.rand()
self.cost(u, s, vt, x)
def test_gradient(self):
(u, s, vt), x = self.manifold.rand()
gu, gs, gvt, gx = self.gradient(u, s, vt, x)
self.assertEqual(gu.shape, (self.m, self.rank))
self.assertEqual(gs.shape, (self.rank, ))
self.assertEqual(gvt.shape, (self.rank, self.n))
self.assertEqual(gx.shape, (self.n, ))
def test_hessian_vector_product(self):
(u, s, vt), x = self.manifold.rand()
(a, b, c), d = self.manifold.rand()
hu, hs, hvt, hx = self.hvp(u, s, vt, x, a, b, c, d)
self.assertEqual(hu.shape, (self.m, self.rank))
self.assertEqual(hs.shape, (self.rank, ))
self.assertEqual(hvt.shape, (self.rank, self.n))
self.assertEqual(hx.shape, (self.n, ))
def test_problem_cost(self):
cost = self.problem.cost
X = self.manifold.rand()
(u, s, vt), x = X
np_testing.assert_allclose(cost(X), self.cost(u, s, vt, x))
def test_problem_egrad(self):
egrad = self.problem.egrad
X = self.manifold.rand()
(u, s, vt), x = X
G = egrad(X)
(gu, gs, gvt), gx = G
for ga, gb in zip((gu, gs, gvt, gx), self.gradient(u, s, vt, x)):
np_testing.assert_allclose(ga, gb)
def test_problem_hessian_vector_product(self):
ehess = self.problem.ehess
X = self.manifold.rand()
U = self.manifold.rand()
H = ehess(X, U)
(u, s, vt), x = X
(a, b, c), d = U
(hu, hs, hvt), hx = H
for ha, hb in zip((hu, hs, hvt, hx), self.hvp(u, s, vt, x, a, b, c,
d)):
np_testing.assert_allclose(ha, hb)
class TestProductManifold(unittest.TestCase):
def setUp(self):
self.m = m = 100
self.n = n = 50
self.euclidean = Euclidean(m, n)
self.sphere = Sphere(n)
self.man = Product([self.euclidean, self.sphere])
def test_dim(self):
np_testing.assert_equal(self.man.dim, self.m*self.n+self.n-1)
def test_typicaldist(self):
np_testing.assert_equal(self.man.typicaldist,
np.sqrt((self.m*self.n)+np.pi**2))
def test_dist(self):
X = self.man.rand()
Y = self.man.rand()
np_testing.assert_equal(self.man.dist(X, Y),
np.sqrt(
self.euclidean.dist(X[0], Y[0])**2 +
self.sphere.dist(X[1], Y[1])**2))
# def test_inner(self):
# def test_proj(self):
# def test_ehess2rhess(self):
# def test_retr(self):
# def test_egrad2rgrad(self):
# def test_norm(self):
# def test_rand(self):
# def test_randvec(self):
# def test_transp(self):
def test_exp_log_inverse(self):
s = self.man
X = s.rand()
Y = s.rand()
Yexplog = s.exp(X, s.log(X, Y))
np_testing.assert_almost_equal(s.dist(Y, Yexplog), 0)
def test_log_exp_inverse(self):
s = self.man
X = s.rand()
U = s.randvec(X)
Ulogexp = s.log(X, s.exp(X, U))
np_testing.assert_array_almost_equal(U[0], Ulogexp[0])
np_testing.assert_array_almost_equal(U[1], Ulogexp[1])
def test_pairmean(self):
s = self.man
X = s.rand()
Y = s.rand()
Z = s.pairmean(X, Y)
np_testing.assert_array_almost_equal(s.dist(X, Z), s.dist(Y, Z))
class TestProductManifold(unittest.TestCase):
def setUp(self):
self.m = m = 100
self.n = n = 50
self.euclidean = Euclidean(m, n)
self.sphere = Sphere(n)
self.man = Product([self.euclidean, self.sphere])
def test_dim(self):
np_testing.assert_equal(self.man.dim, self.m * self.n + self.n - 1)
def test_typicaldist(self):
np_testing.assert_equal(self.man.typicaldist,
np.sqrt((self.m * self.n) + np.pi**2))
def test_dist(self):
X = self.man.rand()
Y = self.man.rand()
np_testing.assert_equal(
self.man.dist(X, Y),
np.sqrt(
self.euclidean.dist(X[0], Y[0])**2 +
self.sphere.dist(X[1], Y[1])**2))
# def test_inner(self):
# def test_proj(self):
# def test_ehess2rhess(self):
# def test_retr(self):
# def test_egrad2rgrad(self):
# def test_norm(self):
# def test_rand(self):
# def test_randvec(self):
# def test_transp(self):
def test_exp_log_inverse(self):
s = self.man
X = s.rand()
Y = s.rand()
Yexplog = s.exp(X, s.log(X, Y))
np_testing.assert_almost_equal(s.dist(Y, Yexplog), 0)
def test_log_exp_inverse(self):
s = self.man
X = s.rand()
U = s.randvec(X)
Ulogexp = s.log(X, s.exp(X, U))
np_testing.assert_array_almost_equal(U[0], Ulogexp[0])
np_testing.assert_array_almost_equal(U[1], Ulogexp[1])
def test_pairmean(self):
s = self.man
X = s.rand()
Y = s.rand()
Z = s.pairmean(X, Y)
np_testing.assert_array_almost_equal(s.dist(X, Z), s.dist(Y, Z))
def fit_gpytorch_manifold(
mll: MarginalLogLikelihood,
bounds: Optional[ParameterBounds] = None,
solver: Solver = pyman_solvers.ConjugateGradient(maxiter=500),
nb_init_candidates: int = 200,
last_x_as_candidate_prob: float = 0.9,
options: Optional[Dict[str, Any]] = None,
track_iterations: bool = True,
approx_mll: bool = False,
module_to_array_func: TModToArray = module_to_list_of_array,
module_from_array_func: TArrayToMod = set_params_with_list_of_array,
) -> Tuple[MarginalLogLikelihood, Dict[str, Union[
float, List[OptimizationIteration]]]]:
"""
This function fits a gpytorch model by maximizing MLL with a pymanopt optimizer.
The model and likelihood in mll must already be in train mode.
This method requires that the model has `train_inputs` and `train_targets`.
Parameters
----------
:param mll: MarginalLogLikelihood to be maximized.
Optional parameters
-------------------
:param nb_init_candidates: number of random initial candidates for the GP parameters
:param last_x_as_candidate_prob: probability that the last set of parameter is among the initial candidates
:param bounds: A dictionary mapping parameter names to tuples of lower and upper bounds.
:param solver: Pymanopt solver.
:param options: Dictionary of solver options, passed along to scipy.minimize.
:param track_iterations: Track the function values and wall time for each iteration.
:param approx_mll: If True, use gpytorch's approximate MLL computation. This is disabled by default since the
stochasticity is an issue for determistic optimizers). Enabling this is only recommended when working with
large training data sets (n>2000).
Returns
-------
:return: 2-element tuple containing
- MarginalLogLikelihood with parameters optimized in-place.
- Dictionary with the following key/values:
"fopt": Best mll value.
"wall_time": Wall time of fitting.
"iterations": List of OptimizationIteration objects with information on each iteration.
If track_iterations is False, will be empty.
Example:
gp = SingleTaskGP(train_X, train_Y)
mll = ExactMarginalLogLikelihood(gp.likelihood, gp)
mll.train()
fit_gpytorch_scipy(mll)
mll.eval()
"""
options = options or {}
# Current parameters
x0, property_dict, bounds = module_to_array_func(module=mll,
bounds=bounds,
exclude=options.pop(
"exclude", None))
x0 = [x0i.astype(np.float64) for x0i in x0]
if bounds is not None:
warnings.warn(
'Bounds handling not supported yet in fit_gpytorch_manifold')
# bounds = Bounds(lb=bounds[0], ub=bounds[1], keep_feasible=True)
t1 = time.time()
# Define cost function
def cost(x):
param_dict = OrderedDict(mll.named_parameters())
idx = 0
for p_name, attrs in property_dict.items():
# Construct the new tensor
if len(attrs.shape) == 0: # deal with scalar tensors
# new_data = torch.tensor(x[0], dtype=attrs.dtype, device=attrs.device)
new_data = torch.tensor(x[idx][0],
dtype=attrs.dtype,
device=attrs.device)
else:
# new_data = torch.tensor(x, dtype=attrs.dtype, device=attrs.device).view(*attrs.shape)
new_data = torch.tensor(x[idx],
dtype=attrs.dtype,
device=attrs.device).view(*attrs.shape)
param_dict[p_name].data = new_data
idx += 1
# mllx = set_params_with_array(mll, x, property_dict)
train_inputs, train_targets = mll.model.train_inputs, mll.model.train_targets
mll.zero_grad()
output = mll.model(*train_inputs)
args = [output, train_targets] + _get_extra_mll_args(mll)
loss = -mll(*args).sum()
return loss
def egrad(x):
loss = cost(x)
loss.backward()
param_dict = OrderedDict(mll.named_parameters())
grad = []
for p_name in property_dict:
t = param_dict[p_name].grad
if t is None:
# this deals with parameters that do not affect the loss
if len(property_dict[p_name].shape
) > 1 and property_dict[p_name].shape[0] > 1:
# if the variable is a matrix, keep its shape
grad.append(np.zeros(property_dict[p_name].shape))
else:
grad.append(np.zeros(property_dict[p_name].shape))
else:
if t.ndim > 1 and t.shape[
0] > 1: # if the variable is a matrix, keep its shape
grad.append(t.detach().cpu().double().clone().numpy())
else: # Vector case
grad.append(
t.detach().view(-1).cpu().double().clone().numpy())
return grad
# Define the manifold (product of manifolds)
manifolds_list = []
for p_name, t in mll.named_parameters():
try:
# If a manifold is given add it
manifolds_list.append(attrgetter(p_name + "_manifold")(mll))
except AttributeError:
# Otherwise, default: Euclidean
manifolds_list.append(
Euclidean(int(np.prod(property_dict[p_name].shape))))
# Product of manifolds
manifold = Product(manifolds_list)
# Instanciate the problem on the manifold
if track_iterations:
verbosity = 2
else:
verbosity = 0
problem = Problem(manifold=manifold,
cost=cost,
egrad=egrad,
verbosity=verbosity,
arg=torch.Tensor()) #, precon=precon)
# For cases where the Hessian is hard/long to compute, we approximate it with finite differences of the gradient.
# Typical cases: the Hessian can be hard to compute due to the 2nd derivative of the eigenvalue decomposition,
# e.g. in the SPD affine-invariant distance.
problem._hess = types.MethodType(get_hessianfd, problem)
# Choose initial parameters
# Do not always consider x0, to encourage variations of the parameters.
if np.random.rand() < last_x_as_candidate_prob:
x0_candidates = [x0]
x0_candidates += [
manifold.rand() for i in range(nb_init_candidates - 1)
]
else:
x0_candidates = []
x0_candidates += [manifold.rand() for i in range(nb_init_candidates)]
for i in range(int(3 * nb_init_candidates / 4)):
x0_candidates[i][0:4] = x0[0:4] #TODO remove hard-coding
y0_candidates = [cost(x0_candidates[i]) for i in range(nb_init_candidates)]
y_init, x_init_idx = torch.Tensor(y0_candidates).min(0)
x_init = x0_candidates[x_init_idx]
with gpt_settings.fast_computations(log_prob=approx_mll):
# Logverbosity of the solver to 1
solver._logverbosity = 1
# Solve
opt_x, opt_log = solver.solve(problem, x=x_init)
# Construct info dict
info_dict = {
"fopt": float(cost(opt_x).detach().numpy()),
"wall_time": time.time() - t1,
"opt_log": opt_log,
}
# if not res.success: # TODO update
# try:
# # Some res.message are bytes
# msg = res.message.decode("ascii")
# except AttributeError:
# # Others are str
# msg = res.message
# warnings.warn(
# f"Fitting failed with the optimizer reporting '{msg}'", OptimizationWarning
# )
# Set to optimum
mll = module_from_array_func(mll, opt_x, property_dict)
return mll, info_dict
class TestProductManifold(unittest.TestCase):
def setUp(self):
self.m = m = 100
self.n = n = 50
self.euclidean = Euclidean(m, n)
self.sphere = Sphere(n)
self.man = Product([self.euclidean, self.sphere])
def test_dim(self):
np_testing.assert_equal(self.man.dim, self.m * self.n + self.n - 1)
def test_typicaldist(self):
np_testing.assert_equal(self.man.typicaldist,
np.sqrt((self.m * self.n) + np.pi**2))
def test_dist(self):
X = self.man.rand()
Y = self.man.rand()
np_testing.assert_equal(
self.man.dist(X, Y),
np.sqrt(
self.euclidean.dist(X[0], Y[0])**2 +
self.sphere.dist(X[1], Y[1])**2))
def test_tangent_vector_multiplication(self):
# Regression test for https://github.com/pymanopt/pymanopt/issues/49.
man = Product((Euclidean(12), Grassmann(12, 3)))
x = man.rand()
eta = man.randvec(x)
np.float64(1.0) * eta
# def test_inner(self):
# def test_proj(self):
# def test_ehess2rhess(self):
# def test_retr(self):
# def test_egrad2rgrad(self):
# def test_norm(self):
# def test_rand(self):
# def test_randvec(self):
# def test_transp(self):
def test_exp_log_inverse(self):
s = self.man
X = s.rand()
Y = s.rand()
Yexplog = s.exp(X, s.log(X, Y))
np_testing.assert_almost_equal(s.dist(Y, Yexplog), 0)
def test_log_exp_inverse(self):
s = self.man
X = s.rand()
U = s.randvec(X)
Ulogexp = s.log(X, s.exp(X, U))
np_testing.assert_array_almost_equal(U[0], Ulogexp[0])
np_testing.assert_array_almost_equal(U[1], Ulogexp[1])
def test_pairmean(self):
s = self.man
X = s.rand()
Y = s.rand()
Z = s.pairmean(X, Y)
np_testing.assert_array_almost_equal(s.dist(X, Z), s.dist(Y, Z))
def test_tangent_vector_multiplication(self):
# Regression test for https://github.com/pymanopt/pymanopt/issues/49.
man = Product((Euclidean(12), Grassmann(12, 3)))
x = man.rand()
eta = man.randvec(x)
np.float64(1.0) * eta
