def loss(self, data, split="train"):
if split == "train":
mask = data.train_mask
elif split == "val":
mask = data.val_mask
else:
mask = data.test_mask
edge_index, edge_types = data.edge_index[:, mask], data.edge_attr[mask]
self.get_edge_set(edge_index, edge_types)
batch_edges, batch_attr, samples, rels, labels = sampling_edge_uniform(
edge_index,
edge_types,
self.edge_set,
self.sampling_rate,
self.num_rels,
label_smoothing=self.lbl_smooth,
num_entities=self.num_entities)
node_embed, rel_embed = self.forward(batch_edges, batch_attr)
sampled_nodes, reindexed_edges = torch.unique(samples,
sorted=True,
return_inverse=True)
assert (self.cache_index == sampled_nodes).any()
loss_n = self._loss(node_embed[reindexed_edges[0]],
node_embed[reindexed_edges[1]], rel_embed[rels],
labels)
loss_r = self.penalty * self._regularization(
[self.emb(sampled_nodes), rel_embed])
return loss_n + loss_r
def loss(self, data, split="train"):
if split == "train":
mask = data.train_mask
elif split == "val":
mask = data.val_mask
else:
mask = data.test_mask
edge_index, edge_types = data.edge_index[:, mask], data.edge_attr[mask]
self.get_edge_set(edge_index, edge_types)
batch_edges, batch_attr, samples, rels, labels = sampling_edge_uniform(
edge_index, edge_types, self.edge_set, self.sampling_rate,
self.num_rels)
output = self.forward(batch_edges, batch_attr)
edge_weight = self.rel_weight(rels)
sampled_nodes, reindexed_edges = torch.unique(samples,
sorted=True,
return_inverse=True)
assert (sampled_nodes == self.cahce_index).any()
sampled_types = torch.unique(rels)
loss_n = self._loss(
output[reindexed_edges[0]], output[reindexed_edges[1]],
edge_weight, labels) + self.penalty * self._regularization(
[self.emb(sampled_nodes),
self.rel_weight(sampled_types)])
return loss_n
def loss(self, data: Graph, split="train"):
if split == "train":
mask = data.train_mask
elif split == "val":
mask = data.val_mask
else:
mask = data.test_mask
row, col = data.edge_index
row, col = row[mask], col[mask]
edge_types = data.edge_attr[mask]
edge_index = torch.stack([row, col])
self.get_edge_set(edge_index, edge_types)
batch_edges, batch_attr, samples, rels, labels = sampling_edge_uniform(
(row, col),
edge_types,
self.edge_set,
self.sampling_rate,
self.num_rels,
label_smoothing=self.lbl_smooth,
num_entities=self.num_entities,
)
with data.local_graph():
data.edge_index = batch_edges
data.edge_attr = batch_attr
node_embed, rel_embed = self.forward(data)
sampled_nodes, reindexed_edges = torch.unique(samples,
sorted=True,
return_inverse=True)
assert (self.cache_index == sampled_nodes).any()
loss_n = self._loss(node_embed[reindexed_edges[0]],
node_embed[reindexed_edges[1]], rel_embed[rels],
labels)
loss_r = self.penalty * self._regularization(
[self.emb(sampled_nodes), rel_embed])
return loss_n + loss_r
