wehub/dmlc--dgl
1
0
Fork 0
Code Issues Pull Requests Actions Deployments Models Agent Notes Packages Projects Releases Wiki Activity

chore: import upstream snapshot with attribution

Browse Source
This commit is contained in:
wehub-resource-sync
2026-07-13 13:35:51 +08:00
commit c36a561cd8
2172 changed files with 455595 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/dgl/nn/pytorch/link/__init__.py
+5
View File
@@ -0,0 +1,5 @@
"""Torch modules for link prediction/knowledge graph completion."""
from .edgepred import EdgePredictor
from .transe import TransE
from .transr import TransR
python/dgl/nn/pytorch/link/edgepred.py
+172
View File
@@ -0,0 +1,172 @@
"""Predictor for edges in homogeneous graphs."""
# pylint: disable= no-member, arguments-differ, invalid-name, W0235
import torch
import torch.nn as nn
import torch.nn.functional as F
class EdgePredictor(nn.Module):
r"""Predictor/score function for pairs of node representations
Given a pair of node representations, :math:`h_i` and :math:`h_j`, it combines them with
**dot product**
.. math::
h_i^{T} h_j
or **cosine similarity**
.. math::
\frac{h_i^{T} h_j}{{\| h_i \|}_2 \cdot {\| h_j \|}_2}
or **elementwise product**
.. math::
h_i \odot h_j
or **concatenation**
.. math::
h_i \Vert h_j
Optionally, it passes the combined results to a linear layer for the final prediction.
Parameters
----------
op : str
The operation to apply. It can be 'dot', 'cos', 'ele', or 'cat',
corresponding to the equations above in order.
in_feats : int, optional
The input feature size of :math:`h_i` and :math:`h_j`. It is required
only if a linear layer is to be applied.
out_feats : int, optional
The output feature size. It is reuiqred only if a linear layer is to be applied.
bias : bool, optional
Whether to use bias for the linear layer if it applies.
Examples
--------
>>> import dgl
>>> import torch as th
>>> from dgl.nn import EdgePredictor
>>> num_nodes = 2
>>> num_edges = 3
>>> in_feats = 4
>>> g = dgl.rand_graph(num_nodes=num_nodes, num_edges=num_edges)
>>> h = th.randn(num_nodes, in_feats)
>>> src, dst = g.edges()
>>> h_src = h[src]
>>> h_dst = h[dst]
Case1: dot product
>>> predictor = EdgePredictor('dot')
>>> predictor(h_src, h_dst).shape
torch.Size([3, 1])
>>> predictor = EdgePredictor('dot', in_feats, out_feats=3)
>>> predictor.reset_parameters()
>>> predictor(h_src, h_dst).shape
torch.Size([3, 3])
Case2: cosine similarity
>>> predictor = EdgePredictor('cos')
>>> predictor(h_src, h_dst).shape
torch.Size([3, 1])
>>> predictor = EdgePredictor('cos', in_feats, out_feats=3)
>>> predictor.reset_parameters()
>>> predictor(h_src, h_dst).shape
torch.Size([3, 3])
Case3: elementwise product
>>> predictor = EdgePredictor('ele')
>>> predictor(h_src, h_dst).shape
torch.Size([3, 4])
>>> predictor = EdgePredictor('ele', in_feats, out_feats=3)
>>> predictor.reset_parameters()
>>> predictor(h_src, h_dst).shape
torch.Size([3, 3])
Case4: concatenation
>>> predictor = EdgePredictor('cat')
>>> predictor(h_src, h_dst).shape
torch.Size([3, 8])
>>> predictor = EdgePredictor('cat', in_feats, out_feats=3)
>>> predictor.reset_parameters()
>>> predictor(h_src, h_dst).shape
torch.Size([3, 3])
"""
def __init__(self, op, in_feats=None, out_feats=None, bias=False):
super(EdgePredictor, self).__init__()
assert op in [
"dot",
"cos",
"ele",
"cat",
], "Expect op to be in ['dot', 'cos', 'ele', 'cat'], got {}".format(op)
self.op = op
if (in_feats is not None) and (out_feats is not None):
if op in ["dot", "cos"]:
in_feats = 1
elif op == "cat":
in_feats = 2 * in_feats
self.linear = nn.Linear(in_feats, out_feats, bias=bias)
else:
self.linear = None
def reset_parameters(self):
r"""
Description
-----------
Reinitialize learnable parameters.
"""
if self.linear is not None:
self.linear.reset_parameters()
def forward(self, h_src, h_dst):
r"""
Description
-----------
Predict for pairs of node representations.
Parameters
----------
h_src : torch.Tensor
Source node features. The tensor is of shape :math:`(E, D_{in})`,
where :math:`E` is the number of edges/node pairs, and :math:`D_{in}`
is the input feature size.
h_dst : torch.Tensor
Destination node features. The tensor is of shape :math:`(E, D_{in})`,
where :math:`E` is the number of edges/node pairs, and :math:`D_{in}`
is the input feature size.
Returns
-------
torch.Tensor
The output features.
"""
if self.op == "dot":
N, D = h_src.shape
h = torch.bmm(h_src.view(N, 1, D), h_dst.view(N, D, 1)).squeeze(-1)
elif self.op == "cos":
h = F.cosine_similarity(h_src, h_dst).unsqueeze(-1)
elif self.op == "ele":
h = h_src * h_dst
else:
h = torch.cat([h_src, h_dst], dim=-1)
if self.linear is not None:
h = self.linear(h)
return h
python/dgl/nn/pytorch/link/transe.py
+99
View File
@@ -0,0 +1,99 @@
"""TransE."""
# pylint: disable= no-member, arguments-differ, invalid-name, W0235
import torch
import torch.nn as nn
class TransE(nn.Module):
r"""Similarity measure from `Translating Embeddings for Modeling Multi-relational Data
<https://papers.nips.cc/paper/2013/hash/1cecc7a77928ca8133fa24680a88d2f9-Abstract.html>`__
Mathematically, it is defined as follows:
.. math::
- {\| h + r - t \|}_p
where :math:`h` is the head embedding, :math:`r` is the relation embedding, and
:math:`t` is the tail embedding.
Parameters
----------
num_rels : int
Number of relation types.
feats : int
Embedding size.
p : int, optional
The p to use for Lp norm, which can be 1 or 2.
Attributes
----------
rel_emb : torch.nn.Embedding
The learnable relation type embedding.
Examples
--------
>>> import dgl
>>> import torch as th
>>> from dgl.nn import TransE
>>> # input features
>>> num_nodes = 10
>>> num_edges = 30
>>> num_rels = 3
>>> feats = 4
>>> scorer = TransE(num_rels=num_rels, feats=feats)
>>> g = dgl.rand_graph(num_nodes=num_nodes, num_edges=num_edges)
>>> src, dst = g.edges()
>>> h = th.randn(num_nodes, feats)
>>> h_head = h[src]
>>> h_tail = h[dst]
>>> # Randomly initialize edge relation types for demonstration
>>> rels = th.randint(low=0, high=num_rels, size=(num_edges,))
>>> scorer(h_head, h_tail, rels).shape
torch.Size([30])
"""
def __init__(self, num_rels, feats, p=1):
super(TransE, self).__init__()
self.rel_emb = nn.Embedding(num_rels, feats)
self.p = p
def reset_parameters(self):
r"""
Description
-----------
Reinitialize learnable parameters.
"""
self.rel_emb.reset_parameters()
def forward(self, h_head, h_tail, rels):
r"""
Description
-----------
Score triples.
Parameters
----------
h_head : torch.Tensor
Head entity features. The tensor is of shape :math:`(E, D)`, where
:math:`E` is the number of triples, and :math:`D` is the feature size.
h_tail : torch.Tensor
Tail entity features. The tensor is of shape :math:`(E, D)`, where
:math:`E` is the number of triples, and :math:`D` is the feature size.
rels : torch.Tensor
Relation types. It is a LongTensor of shape :math:`(E)`, where
:math:`E` is the number of triples.
Returns
-------
torch.Tensor
The triple scores. The tensor is of shape :math:`(E)`.
"""
h_rel = self.rel_emb(rels)
return -torch.norm(h_head + h_rel - h_tail, p=self.p, dim=-1)
python/dgl/nn/pytorch/link/transr.py
+108
View File
@@ -0,0 +1,108 @@
"""TransR."""
# pylint: disable= no-member, arguments-differ, invalid-name, W0235
import torch
import torch.nn as nn
class TransR(nn.Module):
r"""Similarity measure from
`Learning entity and relation embeddings for knowledge graph completion
<https://ojs.aaai.org/index.php/AAAI/article/view/9491>`__
Mathematically, it is defined as follows:
.. math::
- {\| M_r h + r - M_r t \|}_p
where :math:`M_r` is a relation-specific projection matrix, :math:`h` is the
head embedding, :math:`r` is the relation embedding, and :math:`t` is the tail embedding.
Parameters
----------
num_rels : int
Number of relation types.
rfeats : int
Relation embedding size.
nfeats : int
Entity embedding size.
p : int, optional
The p to use for Lp norm, which can be 1 or 2.
Attributes
----------
rel_emb : torch.nn.Embedding
The learnable relation type embedding.
rel_project : torch.nn.Embedding
The learnable relation-type-specific projection.
Examples
--------
>>> import dgl
>>> import torch as th
>>> from dgl.nn import TransR
>>> # input features
>>> num_nodes = 10
>>> num_edges = 30
>>> num_rels = 3
>>> feats = 4
>>> scorer = TransR(num_rels=num_rels, rfeats=2, nfeats=feats)
>>> g = dgl.rand_graph(num_nodes=num_nodes, num_edges=num_edges)
>>> src, dst = g.edges()
>>> h = th.randn(num_nodes, feats)
>>> h_head = h[src]
>>> h_tail = h[dst]
>>> # Randomly initialize edge relation types for demonstration
>>> rels = th.randint(low=0, high=num_rels, size=(num_edges,))
>>> scorer(h_head, h_tail, rels).shape
torch.Size([30])
"""
def __init__(self, num_rels, rfeats, nfeats, p=1):
super(TransR, self).__init__()
self.rel_emb = nn.Embedding(num_rels, rfeats)
self.rel_project = nn.Embedding(num_rels, nfeats * rfeats)
self.rfeats = rfeats
self.nfeats = nfeats
self.p = p
def reset_parameters(self):
r"""
Description
-----------
Reinitialize learnable parameters.
"""
self.rel_emb.reset_parameters()
self.rel_project.reset_parameters()
def forward(self, h_head, h_tail, rels):
r"""
Score triples.
Parameters
----------
h_head : torch.Tensor
Head entity features. The tensor is of shape :math:`(E, D)`, where
:math:`E` is the number of triples, and :math:`D` is the feature size.
h_tail : torch.Tensor
Tail entity features. The tensor is of shape :math:`(E, D)`, where
:math:`E` is the number of triples, and :math:`D` is the feature size.
rels : torch.Tensor
Relation types. It is a LongTensor of shape :math:`(E)`, where
:math:`E` is the number of triples.
Returns
-------
torch.Tensor
The triple scores. The tensor is of shape :math:`(E)`.
"""
h_rel = self.rel_emb(rels)
proj_rel = self.rel_project(rels).reshape(-1, self.nfeats, self.rfeats)
h_head = (h_head.unsqueeze(1) @ proj_rel).squeeze(1)
h_tail = (h_tail.unsqueeze(1) @ proj_rel).squeeze(1)
return -torch.norm(h_head + h_rel - h_tail, p=self.p, dim=-1)