def __init__(self, input_features, n_classes, manifold=None,
bias=True, left_addition=True, signed=True):
super(MobiusClassificationLayer, self).__init__()
self.manifold = Mobius() if manifold is None else manifold
self.input_features = input_features
self.n_classes = n_classes
self.left_addition = left_addition
self.signed = signed
# Initiallize
self.weight = nn.Parameter(torch.Tensor(n_classes, input_features))
self.bias = nn.Parameter(torch.Tensor(n_classes, input_features))
xavier_uniform_(self.weight)
self.bias.data.uniform_(-.01, .01)
def __init__(self,
reg_ot=1e-1,
ot_solver='sinkhorn_knopp',
batch_size=128,
wrapped_function=None,
normalization='max',
is_hyperbolic=True,
match_targets=False):
self.reg_ot = reg_ot
self.ot_solver = ot_solver
self.batch_size = batch_size
self.wrapped_function = wrapped_function
self.normalization = normalization
self.is_hyperbolic = is_hyperbolic
self.match_targets = match_targets
self.manifold = Mobius() if is_hyperbolic else Euclidean()
def __init__(self, input_features, output_features, manifold=None,
bias=True, left_addition=True):
super(MobiusLinear, self).__init__()
self.manifold = Mobius() if manifold is None else manifold
self.input_features = input_features
self.output_features = output_features
self.left_addition = left_addition
# Initiallize
self.weight = nn.Parameter(torch.Tensor(output_features,
input_features))
if bias:
self.bias = nn.Parameter(torch.Tensor(output_features))
else:
self.register_parameter('bias', None)
xavier_uniform_(self.weight)
if self.bias is not None:
self.bias.data.uniform_(-.01, .01)
class MobiusClassificationLayer(nn.Module):
def __init__(self, input_features, n_classes, manifold=None,
bias=True, left_addition=True, signed=True):
super(MobiusClassificationLayer, self).__init__()
self.manifold = Mobius() if manifold is None else manifold
self.input_features = input_features
self.n_classes = n_classes
self.left_addition = left_addition
self.signed = signed
# Initiallize
self.weight = nn.Parameter(torch.Tensor(n_classes, input_features))
self.bias = nn.Parameter(torch.Tensor(n_classes, input_features))
xavier_uniform_(self.weight)
self.bias.data.uniform_(-.01, .01)
def forward(self, x):
output = []
for c in range(self.n_classes):
a = self.weight[c, :]
p = self.bias[c, :]
norm_a_p = self.manifold.norm(p, a) #* self.manifold.s
decision = self.manifold.dist2plane(
x=x, p=p, a=a, signed=self.signed) * norm_a_p
output.append(decision)
return torch.cat(output, dim=1)
def optim_params(self):
""" To use with GeoOpt
"""
return [{
'params': self.parameters(),
'rgrad': self.manifold.rgrad,
'expm': self.manifold.expm,
'logm': self.manifold.logm,
}, ]
def compute_cost(Xs,
Xt,
ys=None,
yt=None,
manifold=None,
is_hyperbolic=True,
normalization='max',
wrapped_function=None,
detach_x=False,
detach_y=True,
limit_max=1e15,
match_targets=False):
"""
Cost used in the OT problem
"""
if is_hyperbolic:
wrapped_function = ((lambda x: torch.cosh(x).log())
if wrapped_function is None else wrapped_function)
manifold = Mobius() if manifold is None else manifold
M_0 = cost_matrix_hyperbolic(Xs,
Xt,
manifold=manifold,
detach_x=detach_x,
detach_y=detach_y)
M_0 = wrapped_function(M_0)
else:
M_0 = cost_matrix(Xs, Xt)
M_0 = cost_normalization(M_0, normalization)
if ((ys is not None) and (yt is not None)) and match_targets:
limit_max_ = limit_max * M_0.max()
ys_ = ys.data.numpy()
yt_ = yt.data.numpy()
classes = [c for c in np.unique(ys_) if c != -1]
# assumes labeled source samples occupy the first rows
# and labeled target samples occupy the first columns
for c in classes:
idx_s = np.where((ys_ != c) & (ys_ != -1))
idx_t = np.where(yt_ == c)
# all the coefficients corresponding to a source sample
# and a target sample :
# with different labels get a infinite
for j in idx_t[0]:
M_0[idx_s[0], j] = torch.tensor([limit_max_])
return M_0
def cost_matrix_hyperbolic(x, y, manifold=None, detach_x=True, detach_y=True):
u, v = x, y
manifold = Mobius() if manifold is not None else manifold
## Hack
try:
if detach_x:
u.detach_()
except:
pass
try:
if detach_y:
v.detach_()
except:
pass
s = manifold.s
#M = acosh(base_poincare_matrix(u, v, s=s))
M = poincare_matrix(u, v, manifold=manifold, s=s)
return M
class MobiusLinear(nn.Module):
def __init__(self, input_features, output_features, manifold=None,
bias=True, left_addition=True):
super(MobiusLinear, self).__init__()
self.manifold = Mobius() if manifold is None else manifold
self.input_features = input_features
self.output_features = output_features
self.left_addition = left_addition
# Initiallize
self.weight = nn.Parameter(torch.Tensor(output_features,
input_features))
if bias:
self.bias = nn.Parameter(torch.Tensor(output_features))
else:
self.register_parameter('bias', None)
xavier_uniform_(self.weight)
if self.bias is not None:
self.bias.data.uniform_(-.01, .01)
def forward(self, x):
if self.bias is not None:
if not self.left_addition:
h = self.manifold.add(
self.manifold.mat_mul(self.weight, x),
self.manifold.expm_zero(self.bias)
)
else:
h = self.manifold.add(
self.manifold.expm_zero(self.bias),
self.manifold.mat_mul(self.weight, x),
)
else:
h = self.manifold.mat_mul(self.weight, x)
return h
def optim_params(self):
""" To use with GeoOpt
"""
return [{
'params': self.parameters(),
'rgrad': self.manifold.rgrad,
'expm': self.manifold.expm,
'logm': self.manifold.logm,
}, ]
def minimize_sinkhorn(Xs,
Xt,
model,
ys=None,
yt=None,
lr=1e-2,
reg_ot=1e-2,
is_hyperbolic=True,
match_targets=True,
n_iter_sinkhorn=100,
max_iter=10000,
stop_thr=1e-5,
max_iter_map=1000,
is_sinkhorn_normalized=False,
every_n=100,
rate=1e-2,
is_projected_output=False,
type_ini='bary',
is_init_map=False,
verbose=False):
mobius = Mobius()
# Initialize: Barycenter approximation
if is_init_map:
model = deepcopy(model)
if is_hyperbolic:
optimizer_init = RiemannianAdam(model.parameters(), lr=lr)
manifold = Mobius()
else:
optimizer_init = Adam(model.parameters(), lr=lr)
manifold = Euclidean()
### Compute cost
_, _, _, coupling = compute_transport(Xs=Xs,
Xt=Xt,
ys=ys,
yt=yt,
reg_ot=reg_ot,
match_targets=match_targets,
manifold=manifold,
is_hyperbolic=is_hyperbolic)
if type_ini == 'bary':
x_approx = manifold.barycenter_mapping(Xt, coupling)
elif (type_ini == 'rot_s2t') or (type_ini == 'rot_t2s'):
xs_np = Xs.data.numpy()
xt_np = Xt.data.numpy()
xs_mean = xs_np.mean(0)
xt_mean = xt_np.mean(0)
xs_centered = xs_np - xs_mean
xt_centered = xt_np - xt_mean
if type_ini == 'rot_s2t':
P, _ = orthogonal_procrustes(xs_centered, xt_centered)
x_approx = torch.FloatTensor(xs_centered.dot(P) + xt_mean)
else:
P, _ = orthogonal_procrustes(xt_centered, xs_centered)
x_approx = torch.FloatTensor(xt_centered.dot(P) + xs_mean)
elif type_ini == 'id':
x_approx = Xs
loop_map = 1 if max_iter_map > 0 else 0
vloss_map = [stop_thr]
it = 0
while loop_map:
it += 1
optimizer_init.zero_grad()
X_pred = mobius.proj2ball(
model(Xs)) if is_projected_output else model(Xs)
loss_map = manifold.distance(X_pred, x_approx).mean()
vloss_map.append(loss_map.item())
relative_error = abs(vloss_map[-1] - vloss_map[-2]) / abs(
vloss_map[-2])
if (it >= max_iter_map) or (np.isnan(vloss_map[-1])):
loop_map = 0
if relative_error < stop_thr:
loop_map = 0
loss_map.backward()
optimizer_init.step()
this_model = deepcopy(model)
lr_mapping = lr * rate
if is_hyperbolic:
optimizer = RiemannianAdam(this_model.parameters(), lr=lr_mapping)
else:
optimizer = Adam(this_model.parameters(), lr=lr_mapping)
vloss = [stop_thr]
loop = 1 if max_iter > 0 else 0
it = 0
while loop:
it += 1
optimizer.zero_grad()
X_pred = mobius.proj2ball(
this_model(Xs)) if is_projected_output else this_model(Xs)
if is_sinkhorn_normalized:
loss = sinkhorn_normalized(X_pred,
Xt,
reg_ot=reg_ot,
n_iter=n_iter_sinkhorn,
ys=ys,
yt=yt,
match_targets=match_targets,
is_hyperbolic=is_hyperbolic)
else:
G, loss = sinkhorn_cost(
X_pred,
Xt,
reg_ot=reg_ot,
n_iter=n_iter_sinkhorn,
match_targets=match_targets,
#wrapped_function=lambda x: -torch.cosh(x),
ys=ys,
yt=yt,
is_hyperbolic=is_hyperbolic)
vloss.append(loss.item())
relative_error = (abs(vloss[-1] - vloss[-2]) /
abs(vloss[-2]) if vloss[-2] != 0 else 0)
if verbose and (it % every_n == 0):
print("\t \t it: %s similarity loss: %.3f" % (it, vloss[-1]))
if (it >= max_iter) or (np.isnan(vloss[-1])):
loop = 0
if relative_error < stop_thr:
loop = 0
loss.backward()
optimizer.step()
return this_model
class HyperbolicSinkhornTransport(BaseEstimator):
def __init__(self,
reg_ot=1e-1,
ot_solver='sinkhorn_knopp',
batch_size=128,
wrapped_function=None,
normalization='max',
is_hyperbolic=True,
match_targets=False):
self.reg_ot = reg_ot
self.ot_solver = ot_solver
self.batch_size = batch_size
self.wrapped_function = wrapped_function
self.normalization = normalization
self.is_hyperbolic = is_hyperbolic
self.match_targets = match_targets
self.manifold = Mobius() if is_hyperbolic else Euclidean()
def fit(self, Xs, Xt, ys=None, yt=None):
self.Xs_train_ = deepcopy(Xs)
self.Xt_train_ = deepcopy(Xt)
_, _, _, coupling = compute_transport(
Xs=Xs,
Xt=Xt,
ys=ys,
yt=yt,
reg_ot=self.reg_ot,
ot_solver=self.ot_solver,
wrapped_function=self.wrapped_function,
normalization=self.normalization,
is_hyperbolic=self.is_hyperbolic,
detach_x=False,
detach_y=False,
match_targets=self.match_targets,
)
self.coupling_ = coupling
return self
def transform(self, Xs):
indices = np.arange(Xs.shape[0])
batch_ind = [
indices[i:i + self.batch_size]
for i in range(0, len(indices), self.batch_size)
]
transp_Xs = []
X_bary = self.manifold.barycenter_mapping(self.Xt_train_,
self.coupling_)
#print(X_bary)
for bi in batch_ind:
# get the nearest neighbor in the source domain
M0 = compute_cost(
Xs[bi],
self.Xs_train_,
normalization=self.normalization,
wrapped_function=lambda x: x, # using the hyperbolic distance
is_hyperbolic=self.is_hyperbolic)
idx = M0.argmin(dim=1)
# define the transported points
diff = self.manifold.add(-self.Xs_train_[idx, :], Xs[bi])
####transp_Xs_ = self.manifold.add(diff, X_bary[idx, :])
transp_Xs_ = self.manifold.add(X_bary[idx, :], diff)
transp_Xs.append(transp_Xs_)
return torch.cat(transp_Xs, dim=0)
def fit_transport(self, Xs, Xt, ys=None, yt=None):
self.fit(Xs=Xs, Xt=Xt, ys=ys, yt=yt)
return self.transform(Xs)
def mapping_estimation(Xs,
Xt,
transport_model=None,
manifold=None,
ys=None,
yt=None,
lr=1e-3,
reg_ot=0,
ot_solver='sinkhorn_knopp',
display_every_map=50,
display_every_coupling=10,
reg_fidelity=1e-3,
stop_thr_coupling=1e-5,
stop_thr_map=1e-5,
stop_thr_global=1e-5,
max_iter_map=300,
max_iter_coupling=20,
wrapped_function=None,
normalization='max',
max_iter=10,
match_targets=False,
is_hyperbolic=True,
verbose=False,
results_path=None,
is_wrapped_linear=False,
save=False):
results_path = '' if results_path is None else results_path
transport_map = deepcopy(transport_model)
if manifold is None:
manifold = Mobius() if is_hyperbolic else Euclidean()
### Compute cost
a, b, M, coupling = compute_transport(Xs=Xs,
Xt=Xt,
ys=ys,
yt=yt,
reg_ot=reg_ot,
ot_solver=ot_solver,
wrapped_function=wrapped_function,
normalization=normalization,
match_targets=match_targets,
is_hyperbolic=is_hyperbolic)
loss = LossModule(xs=Xs,
xt=Xt,
M=M,
a=a,
b=b,
manifold=manifold,
reg_fidelity=reg_fidelity,
ot_solver=ot_solver,
reg_ot=reg_ot)
vloss = [stop_thr_global]
### Initialize transport map
if is_wrapped_linear:
transport_map = solve_wrapped_model(coupling, loss)
else:
transport_map = solve_hyperbolic_nn_model(
transport_map,
coupling,
loss,
max_iter=max_iter_map,
stop_thr=stop_thr_map,
lr=lr,
display_every=display_every_map,
verbose=verbose,
is_hyperbolic=is_hyperbolic)
### Save initial model
if save:
joblib.dump(
transport_map,
pjoin(
results_path,
f"{reg_ot}_{reg_fidelity}_{is_hyperbolic}_{match_targets}_initial.model"
))
vloss.append(to_float(loss.cost(transport_map, coupling)))
loop = 1 if max_iter > 0 else 0
it = 0
if verbose:
print("global it: %s cost: %.4f" % (it, vloss[-1]))
#### Block coordinate descent
while loop:
it += 1
# update coupling
if verbose:
print("\t update coupling")
coupling = solve_coupling(transport_map,
coupling,
loss=loss,
max_iter=max_iter_coupling,
stop_thr=stop_thr_coupling,
verbose=verbose,
every_n=display_every_coupling)
# update transport map
if verbose:
print("\t update transport map")
if is_wrapped_linear:
transport_map = solve_wrapped_model(coupling, loss)
else:
transport_map = solve_hyperbolic_nn_model(
transport_map,
coupling,
loss,
max_iter=max_iter_map,
stop_thr=stop_thr_map,
lr=lr,
display_every=display_every_map,
verbose=verbose,
is_hyperbolic=is_hyperbolic)
if save:
joblib.dump(
transport_map,
pjoin(
results_path,
f"{reg_ot}_{reg_fidelity}_{is_hyperbolic}_{match_targets}_{it}.model"
))
cost = to_float(loss.cost(transport_map, coupling))
vloss.append(cost)
if verbose:
print("global it: %s cost: %.4f" % (it, vloss[-1]))
if (it >= max_iter) or (np.isnan(vloss[-1])):
loop = 0
relative_error = abs(vloss[-1] - vloss[-2]) / abs(vloss[-2])
if relative_error < stop_thr_global:
loop = 0
return transport_map, coupling
def make_domain_adaptation_toy_data(n=1000,
d=2,
sigma=.01,
manifold=None,
bias_source=None,
bias_target=None,
A=None):
manifold = Mobius(eps=1e-15) if manifold is None else manifold
offset = torch.FloatTensor(np.array([[0, .2]]))
bias_source = torch.FloatTensor(np.array(
[[0., .0]])) if bias_source is None else bias_source
bias_target = torch.FloatTensor(np.array(
[[.4, .2]])) if bias_target is None else bias_target
# Source samples
angles = np.random.rand(n, 1) * 2 * np.pi
xs_base = torch.FloatTensor(0.1 * np.concatenate(
(np.sin(angles), np.cos(angles)), axis=1))
noise = torch.randn(n, 2)
xs = manifold.add(manifold.expm_zero(sigma * noise), xs_base)
xs[:n // 2, :] = manifold.add(offset, xs[:n // 2, :])
xs = manifold.add(bias_source, xs)
# Target samples
anglet = np.random.rand(n, 1) * 2 * np.pi
xt_base = torch.FloatTensor(0.1 * np.concatenate(
(np.sin(anglet), np.cos(anglet)), axis=1))
noise = torch.randn(n, 2)
xt = manifold.add(manifold.expm_zero(sigma * noise), xt_base)
xt[:n // 2, :] = manifold.add(offset, xt[:n // 2, :])
# Affine transformation
A = torch.FloatTensor(np.array([[1.5, .7], [.7, 1.5]])) if A is None else A
xt = manifold.add(bias_target, manifold.mat_mul(A, xt))
return xs, xt, A
def __init__(self, manifold=None, hyperbolic_output=True):
super(MobiusElu, self).__init__()
self.manifold = Mobius() if manifold is None else manifold
self.hyperbolic_output = hyperbolic_output
def forward(self, x):
output = []
for c in range(self.n_classes):
a = self.weight[c, :]
p = self.bias[c, :]
norm_a_p = self.manifold.norm(p, a) #* self.manifold.s
decision = self.manifold.dist2plane(
x=x, p=p, a=a, signed=self.signed) * norm_a_p
output.append(decision)
return torch.cat(output, dim=1)
def optim_params(self):
""" To use with GeoOpt
"""
return [{
'params': self.parameters(),
'rgrad': self.manifold.rgrad,
'expm': self.manifold.expm,
'logm': self.manifold.logm,
}, ]
if __name__ == "__main__":
mobius = Mobius(eps=1e-15)
model = nn.Sequential(*[
MobiusLinear(input_features=2, output_features=2, manifold=mobius),
MobiusRelu(manifold=mobius),
MobiusLinear(input_features=2, output_features=2, manifold=mobius),
])