import os
import sys
import torch
import random
import numpy as np
from torch import nn
from torch import optim
from tqdm import trange, tqdm
from collections import Counter
from datetime import datetime
from tensorboardX import SummaryWriter
from torch.utils.data import Dataset, DataLoader
import matplotlib
matplotlib.use("Agg") # this needs to come before other matplotlib imports
import matplotlib.pyplot as plt
plt.style.use("ggplot")
def arcosh(x):
return torch.log(x + torch.sqrt(x ** 2 - 1))
def lorentz_scalar_product(x, y):
# BD, BD -> B
m = x * y
result = m[:, 1:].sum(dim=1) - m[:, 0]
return result
def tangent_norm(x):
# BD -> B
return torch.sqrt(lorentz_scalar_product(x, x))
def exp_map(x, v):
# BD, BD -> BD
tn = tangent_norm(v).unsqueeze(dim=1)
tn_expand = tn.repeat(1, x.size()[-1])
result = torch.cosh(tn) * x + torch.sinh(tn) * (v / tn)
result = torch.where(tn_expand > 0, result, x) # only update if tangent norm is > 0
return result
def set_dim0(x):
x = torch.renorm(x, p=2, dim=0, maxnorm=1e2) # otherwise leaves will explode
# NOTE: the paper does not mention the square part of the equation but if
# you try to derive it you get a square term in the equation
dim0 = torch.sqrt(1 + (x[:, 1:] ** 2).sum(dim=1))
x[:, 0] = dim0
return x
# ========================= models
class RSGD(optim.Optimizer):
def __init__(self, params, learning_rate=None):
learning_rate = learning_rate if learning_rate is not None else 0.01
defaults = {"learning_rate": learning_rate}
super().__init__(params, defaults=defaults)
def step(self):
for group in self.param_groups:
for p in group["params"]:
if p.grad is None:
continue
B, D = p.size()
gl = torch.eye(D, device=p.device, dtype=p.dtype)
gl[0, 0] = -1
grad_norm = torch.norm(p.grad.data)
grad_norm = torch.where(grad_norm > 1, grad_norm, torch.tensor(1.0))
# only normalize if global grad_norm is more than 1
h = (p.grad.data / grad_norm) @ gl
proj = (
h
- (
lorentz_scalar_product(p, h) / lorentz_scalar_product(p, p)
).unsqueeze(1)
* p
)
# print(p, lorentz_scalar_product(p, p))
update = exp_map(p, -group["learning_rate"] * proj)
is_nan_inf = torch.isnan(update) | torch.isinf(update)
update = torch.where(is_nan_inf, p, update)
update[0, :] = p[0, :] # no ❤️ for embedding
update = set_dim0(update)
p.data.copy_(update)
class Lorentz(nn.Module):
"""
This will embed `n_items` in a `dim` dimensional lorentz space.
"""
def __init__(self, n_items, dim, init_range=0.001):
super().__init__()
self.n_items = n_items
self.dim = dim
self.table = nn.Embedding(n_items + 1, dim, padding_idx=0)
nn.init.uniform_(self.table.weight, -init_range, init_range)
# equation 6
with torch.no_grad():
self.table.weight[0] = 5 # padding idx push it to corner
set_dim0(self.table.weight)
def forward(self, I, Ks):
"""
Using the pairwise similarity matrix, generate the following inputs and
provide to this function.
Inputs:
- I : - long tensor
- size (B,)
- This denotes the `i` used in all equations.
- Ks : - long tensor
- size (B, N)
- This denotes at max `N` documents which come from the
nearest neighbor sample.
- The `j` document must be the first of the N indices.
This is used to calculate the losses
Return:
- size (B,)
- Ranking loss calculated using
document to the given `i` document.
"""
n_ks = Ks.size()[1]
ui = torch.stack([self.table(I)] * n_ks, dim=1)
uks = self.table(Ks)
# ---------- reshape for calculation
B, N, D = ui.size()
ui = ui.reshape(B * N, D)
uks = uks.reshape(B * N, D)
dists = -lorentz_scalar_product(ui, uks)
dists = torch.where(dists <= 1, torch.ones_like(dists) + 1e-6, dists)
# sometimes 2 embedding can come very close in R^D.
# when calculating the lorenrz inner product,
# -1 can become -0.99(no idea!), then arcosh will become nan
dists = -arcosh(dists)
# print(dists)
# ---------- turn back to per-sample shape
dists = dists.reshape(B, N)
loss = -(dists[:, 0] - torch.log(torch.exp(dists).sum(dim=1) + 1e-6))
return loss
def lorentz_to_poincare(self):
table = self.table.weight.data.numpy()
return table[:, 1:] / (
table[:, :1] + 1
) # diffeomorphism transform to poincare ball
def get_lorentz_table(self):
return self.table.weight.data.numpy()
def _test_table(self):
x = self.table.weight.data
check = lorentz_scalar_product(x, x) + 1.0
return check.numpy().sum()
class Graph(Dataset):
def __init__(self, pairwise_matrix, sample_size=10):
self.pairwise_matrix = pairwise_matrix
self.n_items = len(pairwise_matrix)
self.sample_size = sample_size
self.arange = np.arange(0, self.n_items)
def __len__(self):
return self.n_items
def __getitem__(self, i):
I = torch.Tensor([i + 1]).squeeze().long()
has_child = (self.pairwise_matrix[i] > 0).sum()
has_parent = (self.pairwise_matrix[:, i] > 0).sum()
arange = np.random.permutation(self.arange)
if has_parent: # if no child go for parent
for j in arange:
if self.pairwise_matrix[j, i] > 0: # assuming no disconneted nodes
min = self.pairwise_matrix[j, i]
break
elif has_child:
for j in arange:
if self.pairwise_matrix[i, j] > 0: # assuming no self loop
min = self.pairwise_matrix[i, j]
break
else:
raise Exception(f"Node {i} has no parent and no child")
arange = np.random.permutation(self.arange)
if has_child:
indices = [x for x in arange if i != x and self.pairwise_matrix[i, x] < min]
else:
indices = [x for x in arange if i != x and self.pairwise_matrix[x, i] < min]
indices = indices[: self.sample_size]
Ks = ([i + 1 for i in [j] + indices] + [0] * self.sample_size)[
: self.sample_size
]
# print(I, Ks)
return I, torch.Tensor(Ks).long()
def recon(table, pair_mat):
"Reconstruction accuracy"
count = 0
table = torch.tensor(table[1:])
for i in range(1, len(pair_mat)): # 0 padding, 1 root, we leave those two
x = table[i].repeat(len(table)).reshape([len(table), len(table[i])]) # N, D
mask = torch.tensor([0.0] * len(table))
mask[i] = 1
mask = mask * -10000.0
dists = lorentz_scalar_product(x, table) + mask
dists = (
dists.numpy()
) # arccosh is monotonically increasing, so no need of that here
# and no -dist also, as acosh in m i, -acosh(-l(x,y)) is nothing but l(x,y)
# print(dists)
predicted_parent = np.argmax(dists)
actual_parent = np.argmax(pair_mat[:, i])
# print(predicted_parent, actual_parent, i, end="\n\n")
count += actual_parent == predicted_parent
count = count / (len(pair_mat) - 1) * 100
return count
_moon_count = 0
def _moon(loss, phases="🌕🌖🌗🌘🌑🌒🌓🌔"):
global _moon_count
_moon_count += 1
p = phases[_moon_count % 8]
return f"{p} Loss: {float(loss)}"
if __name__ == "__main__":
import argparse
parser = argparse.ArgumentParser()
parser.add_argument("dataset", help="File:pairwise_matrix")
parser.add_argument(
"-sample_size", help="How many samples in the N matrix", default=5, type=int
)
parser.add_argument(
"-batch_size", help="How many samples in the batch", default=32, type=int
)
parser.add_argument(
"-burn_c",
help="Divide learning rate by this for the burn epochs",
default=10,
type=int,
)
parser.add_argument(
"-burn_epochs",
help="How many epochs to run the burn phase for?",
default=100,
type=int,
)
parser.add_argument(
"-plot", help="Plot the embeddings", default=False, action="store_true"
)
parser.add_argument("-plot_size", help="Size of the plot", default=3, type=int)
parser.add_argument(
"-plot_graph",
help="Plot the Graph associated with the embeddings",
default=False,
action="store_true",
)
parser.add_argument(
"-overwrite_plots",
help="Overwrite the plots?",
default=False,
action="store_true",
)
parser.add_argument(
"-ckpt", help="Which checkpoint to use?", default=None, type=str
)
parser.add_argument(
"-shuffle", help="Shuffle within batch while learning?", default=True, type=bool
)
parser.add_argument(
"-epochs", help="How many epochs to optimize for?", default=1_000_000, type=int
)
parser.add_argument(
"-poincare_dim",
help="Poincare projection time. Lorentz will be + 1",
default=2,
type=int,
)
parser.add_argument(
"-n_items", help="How many items to embed?", default=None, type=int
)
parser.add_argument(
"-learning_rate", help="RSGD learning rate", default=0.1, type=float
)
parser.add_argument(
"-log_step", help="Log at what multiple of epochs?", default=1, type=int
)
parser.add_argument(
"-logdir", help="What folder to put logs in", default="runs", type=str
)
parser.add_argument(
"-save_step", help="Save at what multiple of epochs?", default=100, type=int
)
parser.add_argument(
"-savedir", help="What folder to put checkpoints in", default="ckpt", type=str
)
parser.add_argument(
"-loader_workers",
help="How many workers to generate tensors",
default=4,
type=int,
)
args = parser.parse_args()
# ----------------------------------- get the correct matrix
if not os.path.exists(args.logdir):
os.mkdir(args.logdir)
if not os.path.exists(args.savedir):
os.mkdir(args.savedir)
exec(f"from datasets import {args.dataset} as pairwise")
pairwise = pairwise[: args.n_items, : args.n_items]
args.n_items = len(pairwise) if args.n_items is None else args.n_items
print(f"{args.n_items} being embedded")
# ---------------------------------- Generate the proper objects
net = Lorentz(
args.n_items, args.poincare_dim + 1
) # as the paper follows R^(n+1) for this space
if args.plot:
if args.poincare_dim != 2:
print("Only embeddings with `-poincare_dim` = 2 are supported for now.")
sys.exit(1)
if args.ckpt is None:
print("Please provide `-ckpt` when using `-plot`")
sys.exit(1)
if os.path.isdir(args.ckpt):
paths = [
os.path.join(args.ckpt, c)
for c in os.listdir(args.ckpt)
if c.endswith("ckpt")
]
else:
paths = [args.ckpt]
paths = list(sorted(paths))
edges = [
tuple(edge)
for edge in set(
[
frozenset((a + 1, b + 1))
for a, row in enumerate(pairwise > 0)
for b, is_non_zero in enumerate(row)
if is_non_zero
]
)
]
print(len(edges), "nodes")
internal_nodes = set(
node
for node, count in Counter(
[node for edge in edges for node in edge]
).items()
if count > 1
)
edges = np.array([edge for edge in edges if edge[1] in internal_nodes])
print(len(edges), "internal nodes")
for path in tqdm(paths, desc="Plotting"):
save_path = f"{path}.svg"
if os.path.exists(save_path) and not args.overwrite_plots:
continue
net.load_state_dict(torch.load(path))
table = net.lorentz_to_poincare()
# skip padding. plot x y
plt.figure(figsize=(7, 7))
if args.plot_graph:
for edge in edges:
plt.plot(
table[edge, 0],
table[edge, 1],
color="black",
marker="o",
alpha=0.5,
)
else:
plt.scatter(table[1:, 0], table[1:, 1])
plt.title(path)
plt.gca().set_xlim(-1, 1)
plt.gca().set_ylim(-1, 1)
plt.gca().add_artist(plt.Circle((0, 0), 1, fill=False, edgecolor="black"))
plt.savefig(save_path)
plt.close()
sys.exit(0)
dataloader = DataLoader(
Graph(pairwise, args.sample_size),
shuffle=args.shuffle,
batch_size=args.batch_size,
num_workers=args.loader_workers,
)
rsgd = RSGD(net.parameters(), learning_rate=args.learning_rate)
name = f"{args.dataset} {datetime.utcnow()}"
writer = SummaryWriter(f"{args.logdir}/{name}")
with tqdm(ncols=80, mininterval=0.2) as epoch_bar:
for epoch in range(args.epochs):
rsgd.learning_rate = (
args.learning_rate / args.burn_c
if epoch < args.burn_epochs
else args.learning_rate
)
for I, Ks in dataloader:
rsgd.zero_grad()
loss = net(I, Ks).mean()
loss.backward()
rsgd.step()
writer.add_scalar("loss", loss, epoch)
writer.add_scalar(
"recon_preform", recon(net.get_lorentz_table(), pairwise), epoch
)
writer.add_scalar("table_test", net._test_table(), epoch)
if epoch % args.save_step == 0:
torch.save(net.state_dict(), f"{args.savedir}/{epoch} {name}.ckpt")
epoch_bar.set_description(
f"🔥 Burn phase loss: {float(loss)}"
if epoch < args.burn_epochs
else _moon(loss)
)
epoch_bar.update(1)
