class LPModel(BaseModel):
"""
Base model for link prediction task.
"""
def __init__(self, args):
super(LPModel, self).__init__(args)
self.dc = FermiDiracDecoder(r=args.r, t=args.t)
self.nb_false_edges = args.nb_false_edges
self.nb_edges = args.nb_edges
def decode(self, h, idx):
if self.manifold_name == 'Euclidean':
h = self.manifold.normalize(h)
emb_in = h[idx[:, 0], :]
emb_out = h[idx[:, 1], :]
sqdist = self.manifold.sqdist(emb_in, emb_out, self.c)
probs = self.dc.forward(sqdist)
assert not torch.isnan(probs).any()
return probs
def compute_metrics(self, embeddings, data, split):
if split == 'train':
edges_false = data[f'{split}_edges_false'][np.random.randint(
0, self.nb_false_edges, self.nb_edges)]
else:
edges_false = data[f'{split}_edges_false']
pos_scores = self.decode(embeddings, data[f'{split}_edges'])
neg_scores = self.decode(embeddings, edges_false)
assert not torch.isnan(pos_scores).any()
assert not torch.isnan(neg_scores).any()
# pos_scores = self.decode(embeddings, data[f'{split}_edges']).clamp(min=1e-8, max=1.0 - 1e-8)
# neg_scores = self.decode(embeddings, edges_false).clamp(min=1e-8, max=1.0 - 1e-8)
loss = F.binary_cross_entropy(pos_scores, torch.ones_like(pos_scores))
loss += F.binary_cross_entropy(neg_scores,
torch.zeros_like(neg_scores))
if pos_scores.is_cuda:
pos_scores = pos_scores.cpu()
neg_scores = neg_scores.cpu()
labels = [1] * pos_scores.shape[0] + [0] * neg_scores.shape[0]
preds = list(pos_scores.data.numpy()) + list(neg_scores.data.numpy())
roc = roc_auc_score(labels, preds)
ap = average_precision_score(labels, preds)
metrics = {'loss': loss, 'roc': roc, 'ap': ap}
return metrics
def init_metric_dict(self):
return {'roc': -1, 'ap': -1}
def has_improved(self, m1, m2):
return 0.5 * (m1['roc'] + m1['ap']) < 0.5 * (m2['roc'] + m2['ap'])
def reset_parameters(self):
print('start reset encoder')
self.encoder.reset_parameters()
self.reset_c()
class LPModel(BaseModel):
"""
Base model for link prediction task.
"""
def __init__(self, args):
super(LPModel, self).__init__(args)
self.dc = FermiDiracDecoder(r=args.r, t=args.t)
self.nb_false_edges = args.nb_false_edges
self.nb_edges = args.nb_edges
def decode(self, h, idx, split):
emb_in = h[idx[:, 0], :]
emb_out = h[idx[:, 1], :]
sqdist = self.manifold.sqdist(emb_in, emb_out, self.c)
probs = self.dc.forward(sqdist, split)
return probs
def compute_metrics(self, embeddings, data, split):
if split == 'train':
edges_false = data[f'{split}_edges_false'][np.random.randint(
0, self.nb_false_edges, self.nb_edges)]
else:
edges_false = data[f'{split}_edges_false']
pos_scores = self.decode(embeddings, data[f'{split}_edges'], split)
neg_scores = self.decode(embeddings, edges_false, split)
loss = F.binary_cross_entropy(pos_scores, torch.ones_like(pos_scores))
loss += F.binary_cross_entropy(neg_scores,
torch.zeros_like(neg_scores))
if pos_scores.is_cuda:
pos_scores = pos_scores.cpu()
neg_scores = neg_scores.cpu()
labels = [1] * pos_scores.shape[0] + [0] * neg_scores.shape[0]
preds = list(pos_scores.data.numpy()) + list(neg_scores.data.numpy())
roc = roc_auc_score(labels, preds)
ap = average_precision_score(labels, preds)
metrics = {'loss': loss, 'roc': roc, 'ap': ap}
return metrics
def init_metric_dict(self):
return {'roc': -1, 'ap': -1}
def has_improved(self, m1, m2):
return 0.5 * (m1['roc'] + m1['ap']) < 0.5 * (m2['roc'] + m2['ap'])
class RELPModel(BaseModel):
"""
Base model for link prediction task.
"""
def __init__(self, args):
super(RELPModel, self).__init__(args)
self.dc = FermiDiracDecoder(r=args.r, t=args.t)
self.nb_false_edges = args.nb_false_edges
self.nb_edges = args.nb_edges
self.net = subnet(32, 32).cuda()
if args.task == 'lp':
args.task = 'nc'
else:
args.task = 'lp'
self.params = torch.load(args.task + 'encoder.pth')
# print(self.params)
self.encoder = getattr(encoders, args.model)(self.c, args)
self.encoder.load_state_dict(self.params)
if args.task == 'lp':
args.task = 'nc'
else:
args.task = 'lp'
self.mark = 0
if args.model == 'GCN' or args.model == 'GAT':
self.mark = 1
# print('self_encoder',self.encoder)
def decode(self, h, idx):
if self.manifold_name == 'Euclidean':
h = self.manifold.normalize(h)
emb_in = h[idx[:, 0], :]
emb_out = h[idx[:, 1], :]
emb = torch.cat((emb_in, emb_out), dim=1)
weights = self.net(emb)
weighed_in = weights[:, :16]
weighed_out = weights[:, 16:]
if self.mark == 1:
sqdist = self.manifold.sqdist(weighed_in, weighed_out, self.c)
else:
sqdist = self.manifold.sqdist(emb_in, emb_out, self.c)
probs = self.dc.forward(sqdist)
return probs
def compute_metrics(self, embeddings, data, split):
if split == 'train':
edges_false = data[f'{split}_edges_false'][np.random.randint(
0, self.nb_false_edges, self.nb_edges)]
else:
edges_false = data[f'{split}_edges_false']
pos_scores = self.decode(embeddings, data[f'{split}_edges'])
neg_scores = self.decode(embeddings, edges_false)
loss = F.binary_cross_entropy(pos_scores, torch.ones_like(pos_scores))
loss += F.binary_cross_entropy(neg_scores,
torch.zeros_like(neg_scores))
if pos_scores.is_cuda:
pos_scores = pos_scores.cpu()
neg_scores = neg_scores.cpu()
labels = [1] * pos_scores.shape[0] + [0] * neg_scores.shape[0]
preds = list(pos_scores.data.numpy()) + list(neg_scores.data.numpy())
preds = np.array(preds)
preds[preds >= 0.5] = 1
preds[preds < 0.5] = 0
preds = list(preds)
print(preds[:5] + preds[-5:])
roc = roc_auc_score(labels, preds)
ap = average_precision_score(labels, preds)
metrics = {'loss': loss, 'roc': roc, 'ap': ap}
return metrics
def init_metric_dict(self):
return {'roc': -1, 'ap': -1}
def has_improved(self, m1, m2):
return 0.5 * (m1['roc'] + m1['ap']) < 0.5 * (m2['roc'] + m2['ap'])
class LPModel(BaseModel):
"""
Base model for link prediction task.
"""
def __init__(self, args):
super(LPModel, self).__init__(args)
self.dc = FermiDiracDecoder(r=args.r, t=args.t)
self.nb_false_edges = args.nb_false_edges
self.nb_edges = args.nb_edges
self.key_param = 'ap'
def decode(self, h, idx):
if self.manifold_name == 'Euclidean':
h = self.manifold.normalize(h)
emb_in = h[idx[:, 0], :]
emb_out = h[idx[:, 1], :]
# decode curvature using last layer's
sqdist = self.manifold.sqdist(emb_in, emb_out, self.c[1])
probs = self.dc.forward(sqdist)
return probs
def compute_metrics(self, embeddings, data, split):
if torch.isnan(embeddings).any():
embeddings[embeddings != embeddings] = 1e-8
if split == 'train':
edges_false = data[f'{split}_edges_false'][np.random.randint(
0, self.nb_false_edges, self.nb_edges)]
else:
edges_false = data[f'{split}_edges_false']
pos_scores = self.decode(embeddings, data[f'{split}_edges'])
neg_scores = self.decode(embeddings, edges_false)
loss = F.binary_cross_entropy(pos_scores, torch.ones_like(pos_scores))
loss += F.binary_cross_entropy(neg_scores,
torch.zeros_like(neg_scores))
if pos_scores.is_cuda:
pos_scores = pos_scores.cpu()
neg_scores = neg_scores.cpu()
labels = [1] * pos_scores.shape[0] + [0] * neg_scores.shape[0]
preds = list(pos_scores.data.numpy()) + list(neg_scores.data.numpy())
roc = roc_auc_score(labels, preds)
ap = average_precision_score(labels, preds)
metrics = {'loss': loss, 'roc': roc, 'ap': ap}
return metrics
def train_with_RL(self, env, Agent1, Agent2, data, epoch, ace):
if epoch >= self.args.start_q:
observation = env.get_observation()
pi1, pi2 = Nash(observation, Agent1, Agent2)
action1 = Agent1.choose_action(observation, pi1)
action2 = Agent2.choose_action(observation, pi2)
actions = (action1, action2)
rewards, train_metrics = env.step((action1, action2), self, data,
ace)
observation_ = env.get_observation()
pi1, pi2 = Nash(observation_, Agent1, Agent2)
pis = (pi1, pi2)
Agent1.learn(observation, actions, rewards[0], observation_, pis)
Agent2.learn(observation, actions, rewards[1], observation_, pis)
if self.args.epsilon_decay == 1:
Agent1.update_epsilon()
Agent2.update_epsilon()
else:
embeddings = self.encode(data['features'], data['adj_train_norm'])
env.embedding1 = env.embedding2 = embeddings
train_metrics = self.compute_metrics(embeddings, data, 'train')
env.acc1_record.append(train_metrics[self.key_param])
env.r1_record.append(0.0)
env.c1_record.append(env.c1[:])
env.acc2_record.append(train_metrics[self.key_param])
env.c2_record.append(env.c1[:])
return train_metrics
def change_curv(self, curv):
'''
Change curvature everywhere in the model.
'''
for item in curv:
if item < 1e-1:
item = 1e-1
self.c = [torch.tensor([c]).float() for c in curv]
if not self.args.cuda == -1:
for i in range(len(self.c)):
self.c[i] = self.c[i].to(self.args.device)
self.encoder.c = self.c
self.encoder.change_curv(self.c)
def init_metric_dict(self):
return {'roc': -1, 'ap': -1}
def has_improved(self, m1, m2):
return 0.5 * (m1['roc'] + m1['ap']) < 0.5 * (m2['roc'] + m2['ap'])
class LPModel(BaseModel):
"""
Base model for link prediction task.
"""
def __init__(self, args):
super(LPModel, self).__init__(args)
self.dc = FermiDiracDecoder(r=args.r, t=args.t)
self.nb_false_edges = args.nb_false_edges
self.nb_edges = args.nb_edges
self.w_e = nn.Linear(args.dim, 1, bias=False)
self.w_h = nn.Linear(args.dim, 1, bias=False)
self.drop_e = args.drop_e
self.drop_h = args.drop_h
self.data = args.dataset
self.model = args.model
self.reset_param()
def reset_param(self):
self.w_e.reset_parameters()
self.w_h.reset_parameters()
def decode(self, h, idx):
if self.manifold_name == 'Euclidean':
h = self.manifold.normalize(h)
if isinstance(h, tuple):
emb_in = h[0][idx[:, 0], :]
emb_out = h[0][idx[:, 1], :]
"compute hyperbolic dist"
emb_in = self.manifold.logmap0(emb_in, self.c)
emb_out = self.manifold.logmap0(emb_out, self.c)
sqdist_h = torch.sqrt((emb_in - emb_out).pow(2).sum(dim=-1) +
1e-15)
probs_h = self.dc.forward(sqdist_h)
"compute dist in Euclidean"
emb_in_e = h[1][idx[:, 0], :]
emb_out_e = h[1][idx[:, 1], :]
sqdist_e = torch.sqrt((emb_in_e - emb_out_e).pow(2).sum(dim=-1) +
1e-15)
probs_e = self.dc.forward(sqdist_e)
# sub
w_h = torch.sigmoid(self.w_h(emb_in - emb_out).view(-1))
w_e = torch.sigmoid(self.w_e(emb_in_e - emb_out_e).view(-1))
w = torch.cat([w_h.view(-1, 1), w_e.view(-1, 1)], dim=-1)
if self.data == 'pubmed':
w = F.normalize(w, p=1, dim=-1)
probs = torch.sigmoid(w[-1, 0] * probs_h + w[-1, 1] * probs_e)
assert torch.min(probs) >= 0
assert torch.max(probs) <= 1
else:
emb_in = h[idx[:, 0], :]
emb_out = h[idx[:, 1], :]
sqdist = self.manifold.sqdist(emb_in, emb_out, self.c)
assert torch.max(sqdist) >= 0
probs = self.dc.forward(sqdist)
return probs
def compute_metrics(self, embeddings, data, split, args):
if split == 'train':
edges_false = data[f'{split}_edges_false'][np.random.randint(
0, self.nb_false_edges, self.nb_edges)]
else:
edges_false = data[f'{split}_edges_false']
pos_scores = self.decode(embeddings, data[f'{split}_edges'])
neg_scores = self.decode(embeddings, edges_false)
loss = F.binary_cross_entropy(pos_scores, torch.ones_like(pos_scores))
loss += F.binary_cross_entropy(neg_scores,
torch.zeros_like(neg_scores))
if pos_scores.is_cuda:
pos_scores = pos_scores.cpu()
neg_scores = neg_scores.cpu()
labels = [1] * pos_scores.shape[0] + [0] * neg_scores.shape[0]
preds = list(pos_scores.data.numpy()) + list(neg_scores.data.numpy())
roc = roc_auc_score(labels, preds)
ap = average_precision_score(labels, preds)
metrics = {'loss': loss, 'roc': roc, 'ap': ap}
return metrics
def init_metric_dict(self):
return {'roc': -1, 'ap': -1}
def has_improved(self, m1, m2):
return 0.5 * (m1['roc'] + m1['ap']) < 0.5 * (m2['roc'] + m2['ap'])