chore: import upstream snapshot with attribution
This commit is contained in:
wehub-resource-sync
@@ -0,0 +1,251 @@
|
||||
"""Heterograph NN modules"""
|
||||
import tensorflow as tf
|
||||
from tensorflow.keras import layers
|
||||
|
||||
__all__ = ["HeteroGraphConv"]
|
||||
|
||||
|
||||
class HeteroGraphConv(layers.Layer):
|
||||
r"""A generic module for computing convolution on heterogeneous graphs.
|
||||
|
||||
The heterograph convolution applies sub-modules on their associating
|
||||
relation graphs, which reads the features from source nodes and writes the
|
||||
updated ones to destination nodes. If multiple relations have the same
|
||||
destination node types, their results are aggregated by the specified method.
|
||||
If the relation graph has no edge, the corresponding module will not be called.
|
||||
|
||||
Pseudo-code:
|
||||
|
||||
.. code::
|
||||
|
||||
outputs = {nty : [] for nty in g.dsttypes}
|
||||
# Apply sub-modules on their associating relation graphs in parallel
|
||||
for relation in g.canonical_etypes:
|
||||
stype, etype, dtype = relation
|
||||
dstdata = relation_submodule(g[relation], ...)
|
||||
outputs[dtype].append(dstdata)
|
||||
|
||||
# Aggregate the results for each destination node type
|
||||
rsts = {}
|
||||
for ntype, ntype_outputs in outputs.items():
|
||||
if len(ntype_outputs) != 0:
|
||||
rsts[ntype] = aggregate(ntype_outputs)
|
||||
return rsts
|
||||
|
||||
Examples
|
||||
--------
|
||||
|
||||
Create a heterograph with three types of relations and nodes.
|
||||
|
||||
>>> import dgl
|
||||
>>> g = dgl.heterograph({
|
||||
... ('user', 'follows', 'user') : edges1,
|
||||
... ('user', 'plays', 'game') : edges2,
|
||||
... ('store', 'sells', 'game') : edges3})
|
||||
|
||||
Create a ``HeteroGraphConv`` that applies different convolution modules to
|
||||
different relations. Note that the modules for ``'follows'`` and ``'plays'``
|
||||
do not share weights.
|
||||
|
||||
>>> import dgl.nn.pytorch as dglnn
|
||||
>>> conv = dglnn.HeteroGraphConv({
|
||||
... 'follows' : dglnn.GraphConv(...),
|
||||
... 'plays' : dglnn.GraphConv(...),
|
||||
... 'sells' : dglnn.SAGEConv(...)},
|
||||
... aggregate='sum')
|
||||
|
||||
Call forward with some ``'user'`` features. This computes new features for both
|
||||
``'user'`` and ``'game'`` nodes.
|
||||
|
||||
>>> import tensorflow as tf
|
||||
>>> h1 = {'user' : tf.random.normal((g.num_nodes('user'), 5))}
|
||||
>>> h2 = conv(g, h1)
|
||||
>>> print(h2.keys())
|
||||
dict_keys(['user', 'game'])
|
||||
|
||||
Call forward with both ``'user'`` and ``'store'`` features. Because both the
|
||||
``'plays'`` and ``'sells'`` relations will update the ``'game'`` features,
|
||||
their results are aggregated by the specified method (i.e., summation here).
|
||||
|
||||
>>> f1 = {'user' : ..., 'store' : ...}
|
||||
>>> f2 = conv(g, f1)
|
||||
>>> print(f2.keys())
|
||||
dict_keys(['user', 'game'])
|
||||
|
||||
Call forward with some ``'store'`` features. This only computes new features
|
||||
for ``'game'`` nodes.
|
||||
|
||||
>>> g1 = {'store' : ...}
|
||||
>>> g2 = conv(g, g1)
|
||||
>>> print(g2.keys())
|
||||
dict_keys(['game'])
|
||||
|
||||
Call forward with a pair of inputs is allowed and each submodule will also
|
||||
be invoked with a pair of inputs.
|
||||
|
||||
>>> x_src = {'user' : ..., 'store' : ...}
|
||||
>>> x_dst = {'user' : ..., 'game' : ...}
|
||||
>>> y_dst = conv(g, (x_src, x_dst))
|
||||
>>> print(y_dst.keys())
|
||||
dict_keys(['user', 'game'])
|
||||
|
||||
Notes
|
||||
-----
|
||||
|
||||
HeteroGraphConv requires that there is a module for every ``'etype'`` in an input graph.
|
||||
If you want to apply HeteroGraphConv to a subset of a graph's ``'etypes'``, you must
|
||||
create a new graph using for example :func:`~dgl.edge_type_subgraph()`.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
mods : dict[str, nn.Module]
|
||||
Modules associated with every edge types. The forward function of each
|
||||
module must have a `DGLGraph` object as the first argument, and
|
||||
its second argument is either a tensor object representing the node
|
||||
features or a pair of tensor object representing the source and destination
|
||||
node features.
|
||||
aggregate : str, callable, optional
|
||||
Method for aggregating node features generated by different relations.
|
||||
Allowed string values are 'sum', 'max', 'min', 'mean', 'stack'.
|
||||
The 'stack' aggregation is performed along the second dimension, whose order
|
||||
is deterministic.
|
||||
User can also customize the aggregator by providing a callable instance.
|
||||
For example, aggregation by summation is equivalent to the follows:
|
||||
|
||||
.. code::
|
||||
|
||||
def my_agg_func(tensors, dsttype):
|
||||
# tensors: is a list of tensors to aggregate
|
||||
# dsttype: string name of the destination node type for which the
|
||||
# aggregation is performed
|
||||
stacked = tf.stack(tensors, axis=0)
|
||||
return tf.reduce_sum(stacked, axis=0)
|
||||
|
||||
Attributes
|
||||
----------
|
||||
mods : dict[str, nn.Module]
|
||||
Modules associated with every edge types.
|
||||
"""
|
||||
|
||||
def __init__(self, mods, aggregate="sum"):
|
||||
super(HeteroGraphConv, self).__init__()
|
||||
self.mods = mods
|
||||
# Do not break if graph has 0-in-degree nodes.
|
||||
# Because there is no general rule to add self-loop for heterograph.
|
||||
for _, v in self.mods.items():
|
||||
set_allow_zero_in_degree_fn = getattr(
|
||||
v, "set_allow_zero_in_degree", None
|
||||
)
|
||||
if callable(set_allow_zero_in_degree_fn):
|
||||
set_allow_zero_in_degree_fn(True)
|
||||
if isinstance(aggregate, str):
|
||||
self.agg_fn = get_aggregate_fn(aggregate)
|
||||
else:
|
||||
self.agg_fn = aggregate
|
||||
|
||||
def call(self, g, inputs, mod_args=None, mod_kwargs=None):
|
||||
"""Forward computation
|
||||
|
||||
Invoke the forward function with each module and aggregate their results.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
g : DGLGraph
|
||||
Graph data.
|
||||
inputs : dict[str, Tensor] or pair of dict[str, Tensor]
|
||||
Input node features.
|
||||
mod_args : dict[str, tuple[any]], optional
|
||||
Extra positional arguments for the sub-modules.
|
||||
mod_kwargs : dict[str, dict[str, any]], optional
|
||||
Extra key-word arguments for the sub-modules.
|
||||
|
||||
Returns
|
||||
-------
|
||||
dict[str, Tensor]
|
||||
Output representations for every types of nodes.
|
||||
"""
|
||||
if mod_args is None:
|
||||
mod_args = {}
|
||||
if mod_kwargs is None:
|
||||
mod_kwargs = {}
|
||||
outputs = {nty: [] for nty in g.dsttypes}
|
||||
if isinstance(inputs, tuple):
|
||||
src_inputs, dst_inputs = inputs
|
||||
for stype, etype, dtype in g.canonical_etypes:
|
||||
rel_graph = g[stype, etype, dtype]
|
||||
if stype not in src_inputs or dtype not in dst_inputs:
|
||||
continue
|
||||
dstdata = self.mods[etype](
|
||||
rel_graph,
|
||||
(src_inputs[stype], dst_inputs[dtype]),
|
||||
*mod_args.get(etype, ()),
|
||||
**mod_kwargs.get(etype, {})
|
||||
)
|
||||
outputs[dtype].append(dstdata)
|
||||
else:
|
||||
for stype, etype, dtype in g.canonical_etypes:
|
||||
rel_graph = g[stype, etype, dtype]
|
||||
if stype not in inputs:
|
||||
continue
|
||||
dstdata = self.mods[etype](
|
||||
rel_graph,
|
||||
(inputs[stype], inputs[dtype]),
|
||||
*mod_args.get(etype, ()),
|
||||
**mod_kwargs.get(etype, {})
|
||||
)
|
||||
outputs[dtype].append(dstdata)
|
||||
rsts = {}
|
||||
for nty, alist in outputs.items():
|
||||
if len(alist) != 0:
|
||||
rsts[nty] = self.agg_fn(alist, nty)
|
||||
return rsts
|
||||
|
||||
|
||||
def get_aggregate_fn(agg):
|
||||
"""Internal function to get the aggregation function for node data
|
||||
generated from different relations.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
agg : str
|
||||
Method for aggregating node features generated by different relations.
|
||||
Allowed values are 'sum', 'max', 'min', 'mean', 'stack'.
|
||||
|
||||
Returns
|
||||
-------
|
||||
callable
|
||||
Aggregator function that takes a list of tensors to aggregate
|
||||
and returns one aggregated tensor.
|
||||
"""
|
||||
if agg == "sum":
|
||||
fn = tf.reduce_sum
|
||||
elif agg == "max":
|
||||
fn = tf.reduce_max
|
||||
elif agg == "min":
|
||||
fn = tf.reduce_min
|
||||
elif agg == "mean":
|
||||
fn = tf.reduce_mean
|
||||
elif agg == "stack":
|
||||
fn = None # will not be called
|
||||
else:
|
||||
raise DGLError(
|
||||
"Invalid cross type aggregator. Must be one of "
|
||||
'"sum", "max", "min", "mean" or "stack". But got "%s"' % agg
|
||||
)
|
||||
if agg == "stack":
|
||||
|
||||
def stack_agg(inputs, dsttype): # pylint: disable=unused-argument
|
||||
if len(inputs) == 0:
|
||||
return None
|
||||
return tf.stack(inputs, axis=1)
|
||||
|
||||
return stack_agg
|
||||
else:
|
||||
|
||||
def aggfn(inputs, dsttype): # pylint: disable=unused-argument
|
||||
if len(inputs) == 0:
|
||||
return None
|
||||
stacked = tf.stack(inputs, axis=0)
|
||||
return fn(stacked, axis=0)
|
||||
|
||||
return aggfn
|
||||
Reference in New Issue
Block a user