def __init__(self,
n_exemple,
size=10,
cuda=False,
lr=1e-2,
verbose=True,
negative_distribution=None,
optimizer_method=optimizer.PoincareBallSGDAdd,
aggregation=torch.sum):
super(PoincareEmbedding, self).__init__()
self.cuda = cuda
self.N = n_exemple
self.W = poincare_module.PoincareEmbedding(n_exemple, size)
self.C = poincare_module.PoincareEmbedding(n_exemple, size)
if (self.cuda):
self.W.cuda()
self.C.cuda()
self.optimizer = optimizer_method(list(self.W.parameters()) +
list(self.C.parameters()),
lr=lr)
self.verbose = verbose
self.lr = lr
self.d = poincare_function.poincare_distance
if (negative_distribution is None):
self.n_dist = torch.distributions.Categorical(
torch.ones(self.N) / self.N)
else:
self.n_dist = negative_distribution
self.agg = aggregation
def __init__(self,
n_exemple,
cuda=False,
lr=1e-2,
verbose=True,
negative_distribution=None):
super(RiemannianEmbedding, self).__init__()
self.cuda = cuda
self.N = n_exemple
self.W = poincare_module.PoincareEmbedding(n_exemple, 2)
if (self.cuda):
self.W.cuda()
self.optimizer = optimizer.PoincareSGDExp(self.W.parameters(), lr=lr)
self.verbose = verbose
self.d = poincare_function.poincare_distance_dg
if (negative_distribution is None):
self.n_dist = torch.distributions.Categorical(
torch.ones(self.N) / self.N)
else:
self.n_dist = negative_distribution