wehub-resource-sync
252 lines
8.7 KiB
Python
252 lines
8.7 KiB
Python
"""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
|