Skip to content

GitLab

Unverified Commit 76bb5404 authored by Hongzhi (Steve), Chen's avatar Hongzhi (Steve), Chen Committed by GitHub
Browse files

[Misc] Black auto fix. (#4682) 


Co-authored-by: default avatarSteve <ubuntu@ip-172-31-34-29.ap-northeast-1.compute.internal>
parent a208e886
Hide whitespace changes
Inline Side-by-side
Showing with 776 additions and 432 deletions
python/dgl/ops/sddmm.py
View file @ 76bb5404
"""dgl sddmm operator module."""
from itertools import product
import sys
from itertools import product
from .. import backend as F
from ..backend import gsddmm as gsddmm_internal
from ..backend import gsddmm_hetero as gsddmm_internal_hetero
from .. import backend as F
__all__ = ['gsddmm', 'copy_u', 'copy_v', 'copy_e']
__all__ = ["gsddmm", "copy_u", "copy_v", "copy_e"]
def reshape_lhs_rhs(lhs_data, rhs_data):
r""" Expand dims so that there will be no broadcasting issues with different
r"""Expand dims so that there will be no broadcasting issues with different
number of dimensions. For example, given two shapes (N, 3, 1), (E, 5, 3, 4)
that are valid broadcastable shapes, change them to (N, 1, 3, 1) and
(E, 5, 3, 4)
......@@ -33,8 +34,9 @@ def reshape_lhs_rhs(lhs_data, rhs_data):
rhs_data = F.reshape(rhs_data, new_rhs_shape)
return lhs_data, rhs_data
def gsddmm(g, op, lhs_data, rhs_data, lhs_target='u', rhs_target='v'):
r""" Generalized Sampled-Dense-Dense Matrix Multiplication interface.
def gsddmm(g, op, lhs_data, rhs_data, lhs_target="u", rhs_target="v"):
r"""Generalized Sampled-Dense-Dense Matrix Multiplication interface.
It computes edge features by :attr:`op` lhs features and rhs features.
.. math::
......@@ -69,30 +71,34 @@ def gsddmm(g, op, lhs_data, rhs_data, lhs_target='u', rhs_target='v'):
The result tensor.
"""
if g._graph.number_of_etypes() == 1:
if op not in ['copy_lhs', 'copy_rhs']:
if op not in ["copy_lhs", "copy_rhs"]:
lhs_data, rhs_data = reshape_lhs_rhs(lhs_data, rhs_data)
return gsddmm_internal(
g._graph, op, lhs_data, rhs_data, lhs_target, rhs_target)
g._graph, op, lhs_data, rhs_data, lhs_target, rhs_target
)
else:
if op == 'copy_lhs':
if op == "copy_lhs":
rhs_data = [None] * g._graph.number_of_etypes()
elif op == 'copy_rhs':
elif op == "copy_rhs":
lhs_data = [None] * g._graph.number_of_ntypes()
# TODO (Israt): Call reshape_lhs_rhs() on lhs and rhs data to match their dimension
# and avoid broadcasting issue. Handle the case where different nodes have
# different dimensions, and different etypes may need different broadcasting
# dims for the same node.
lhs_and_rhs_tuple = tuple(list(lhs_data) + list(rhs_data))
return gsddmm_internal_hetero(g._graph, op, len(lhs_data), lhs_target,
rhs_target, *lhs_and_rhs_tuple)
return gsddmm_internal_hetero(
g._graph,
op,
len(lhs_data),
lhs_target,
rhs_target,
*lhs_and_rhs_tuple
)
def _gen_sddmm_func(lhs_target, rhs_target, binary_op):
name = "{}_{}_{}".format(lhs_target, binary_op, rhs_target)
target_dict = {
'u': "source node",
'e': "edge",
'v': "destination node"
}
target_dict = {"u": "source node", "e": "edge", "v": "destination node"}
lhs_str = target_dict[lhs_target]
rhs_str = target_dict[rhs_target]
docstring = r"""Generalized SDDMM function.
......@@ -121,11 +127,15 @@ def _gen_sddmm_func(lhs_target, rhs_target, binary_op):
Broadcasting follows NumPy semantics. Please see
https://docs.scipy.org/doc/numpy/user/basics.broadcasting.html
for more details about the NumPy broadcasting semantics.
""".format(op=binary_op, lhs=lhs_str, rhs=rhs_str)
""".format(
op=binary_op, lhs=lhs_str, rhs=rhs_str
)
def func(g, x, y):
return gsddmm(g, binary_op, x, y,
lhs_target=lhs_target, rhs_target=rhs_target)
return gsddmm(
g, binary_op, x, y, lhs_target=lhs_target, rhs_target=rhs_target
)
func.__name__ = name
func.__doc__ = docstring
return func
......@@ -161,7 +171,7 @@ def copy_u(g, x):
-----
This function supports autograd (computing input gradients given the output gradient).
"""
return gsddmm(g, 'copy_lhs', x, None)
return gsddmm(g, "copy_lhs", x, None)
def copy_v(g, x):
......@@ -183,7 +193,7 @@ def copy_v(g, x):
-----
This function supports autograd (computing input gradients given the output gradient).
"""
return gsddmm(g, 'copy_rhs', None, x)
return gsddmm(g, "copy_rhs", None, x)
# pylint: disable=unused-argument
......
python/dgl/ops/segment.py
View file @ 76bb5404
"""Segment aggregation operators implemented using DGL graph."""
from ..base import DGLError
from .. import backend as F
from ..base import DGLError
__all__ = ['segment_reduce', 'segment_softmax', 'segment_mm']
__all__ = ["segment_reduce", "segment_softmax", "segment_mm"]
def segment_reduce(seglen, value, reducer='sum'):
def segment_reduce(seglen, value, reducer="sum"):
"""Segment reduction operator.
It aggregates the value tensor along the first dimension by segments.
......@@ -41,16 +42,17 @@ def segment_reduce(seglen, value, reducer='sum'):
[4., 4., 4.]])
"""
offsets = F.cumsum(
F.cat([F.zeros((1,), F.dtype(seglen), F.context(seglen)), seglen], 0), 0)
if reducer == 'mean':
rst = F.segment_reduce('sum', value, offsets)
F.cat([F.zeros((1,), F.dtype(seglen), F.context(seglen)), seglen], 0), 0
)
if reducer == "mean":
rst = F.segment_reduce("sum", value, offsets)
rst_shape = F.shape(rst)
z = F.astype(F.clamp(seglen, 1, len(value)), F.dtype(rst))
z_shape = (rst_shape[0],) + (1,) * (len(rst_shape) - 1)
return rst / F.reshape(z, z_shape)
elif reducer in ['min', 'sum', 'max']:
elif reducer in ["min", "sum", "max"]:
rst = F.segment_reduce(reducer, value, offsets)
if reducer in ['min', 'max']:
if reducer in ["min", "max"]:
rst = F.replace_inf_with_zero(rst)
return rst
else:
......@@ -95,13 +97,14 @@ def segment_softmax(seglen, value):
[0.2500, 0.2500, 0.2500],
[0.2500, 0.2500, 0.2500]])
"""
value_max = segment_reduce(seglen, value, reducer='max')
value_max = segment_reduce(seglen, value, reducer="max")
value = F.exp(value - F.repeat(value_max, seglen, dim=0))
value_sum = segment_reduce(seglen, value, reducer='sum')
value_sum = segment_reduce(seglen, value, reducer="sum")
return value / F.repeat(value_sum, seglen, dim=0)
def segment_mm(a, b, seglen_a):
r""" Performs matrix multiplication according to segments.
r"""Performs matrix multiplication according to segments.
Suppose ``seglen_a == [10, 5, 0, 3]``, the operator will perform
four matrix multiplications::
......
python/dgl/ops/spmm.py
View file @ 76bb5404
"""dgl spmm operator module."""
import sys
from .. import backend as F
from ..backend import gspmm as gspmm_internal
from ..backend import gspmm_hetero as gspmm_internal_hetero
from .. import backend as F
__all__ = ['gspmm']
__all__ = ["gspmm"]
def reshape_lhs_rhs(lhs_data, rhs_data):
r""" Expand dims so that there will be no broadcasting issues with different
r"""Expand dims so that there will be no broadcasting issues with different
number of dimensions. For example, given two shapes (N, 3, 1), (E, 5, 3, 4)
that are valid broadcastable shapes, change them to (N, 1, 3, 1) and
(E, 5, 3, 4)
......@@ -32,8 +33,9 @@ def reshape_lhs_rhs(lhs_data, rhs_data):
rhs_data = F.reshape(rhs_data, new_rhs_shape)
return lhs_data, rhs_data
def gspmm(g, op, reduce_op, lhs_data, rhs_data):
r""" Generalized Sparse Matrix Multiplication interface.
r"""Generalized Sparse Matrix Multiplication interface.
It fuses two steps into one kernel.
1. Computes messages by :attr:`op` source node and edge features.
......@@ -69,27 +71,45 @@ def gspmm(g, op, reduce_op, lhs_data, rhs_data):
The result tensor.
"""
if g._graph.number_of_etypes() == 1:
if op not in ['copy_lhs', 'copy_rhs']:
if op not in ["copy_lhs", "copy_rhs"]:
lhs_data, rhs_data = reshape_lhs_rhs(lhs_data, rhs_data)
# With max and min reducers infinity will be returned for zero degree nodes
ret = gspmm_internal(g._graph, op,
'sum' if reduce_op == 'mean' else reduce_op,
lhs_data, rhs_data)
ret = gspmm_internal(
g._graph,
op,
"sum" if reduce_op == "mean" else reduce_op,
lhs_data,
rhs_data,
)
else:
# lhs_data or rhs_data is None only in unary functions like ``copy-u`` or ``copy_e``
lhs_data = [None] * g._graph.number_of_ntypes() if lhs_data is None else lhs_data
rhs_data = [None] * g._graph.number_of_etypes() if rhs_data is None else rhs_data
lhs_data = (
[None] * g._graph.number_of_ntypes()
if lhs_data is None
else lhs_data
)
rhs_data = (
[None] * g._graph.number_of_etypes()
if rhs_data is None
else rhs_data
)
# TODO (Israt): Call reshape func
lhs_and_rhs_tuple = tuple(list(lhs_data) + list(rhs_data))
ret = gspmm_internal_hetero(g._graph, op,
'sum' if reduce_op == 'mean' else reduce_op,
len(lhs_data), *lhs_and_rhs_tuple)
ret = gspmm_internal_hetero(
g._graph,
op,
"sum" if reduce_op == "mean" else reduce_op,
len(lhs_data),
*lhs_and_rhs_tuple
)
# TODO (Israt): Add support for 'mean' in heterograph
# divide in degrees for mean reducer.
if reduce_op == 'mean':
if reduce_op == "mean":
ret_shape = F.shape(ret)
deg = g.in_degrees()
deg = F.astype(F.clamp(deg, 1, max(g.number_of_edges(), 1)), F.dtype(ret))
deg = F.astype(
F.clamp(deg, 1, max(g.number_of_edges(), 1)), F.dtype(ret)
)
deg_shape = (ret_shape[0],) + (1,) * (len(ret_shape) - 1)
return ret / F.reshape(deg, deg_shape)
else:
......@@ -98,13 +118,17 @@ def gspmm(g, op, reduce_op, lhs_data, rhs_data):
def _attach_zerodeg_note(docstring, reducer):
note1 = """
The {} function will return zero for nodes with no incoming messages.""".format(reducer)
The {} function will return zero for nodes with no incoming messages.""".format(
reducer
)
note2 = """
This is implemented by replacing all {} values to zero.
""".format("infinity" if reducer == "min" else "negative infinity")
""".format(
"infinity" if reducer == "min" else "negative infinity"
)
docstring = docstring + note1
if reducer in ('min', 'max'):
if reducer in ("min", "max"):
docstring = docstring + note2
return docstring
......@@ -140,11 +164,14 @@ def _gen_spmm_func(binary_op, reduce_op):
Broadcasting follows NumPy semantics. Please see
https://docs.scipy.org/doc/numpy/user/basics.broadcasting.html
for more details about the NumPy broadcasting semantics.
""".format(binary_op, reduce_op)
""".format(
binary_op, reduce_op
)
docstring = _attach_zerodeg_note(docstring, reduce_op)
def func(g, x, y):
return gspmm(g, binary_op, reduce_op, x, y)
func.__name__ = name
func.__doc__ = docstring
return func
......@@ -155,13 +182,11 @@ def _gen_copy_reduce_func(binary_op, reduce_op):
name = "{}_{}".format(binary_op, reduce_op)
binary_str = {
"copy_u": "It copies node feature to edge as the message.",
'copy_e': "It regards edge feature as message."
}
x_str = {
"copy_u": "source node",
"copy_e": "edge"
"copy_e": "It regards edge feature as message.",
}
docstring = lambda binary_op: _attach_zerodeg_note("""Generalized SpMM function. {}
x_str = {"copy_u": "source node", "copy_e": "edge"}
docstring = lambda binary_op: _attach_zerodeg_note(
"""Generalized SpMM function. {}
Then aggregates the message by {} on destination nodes.
Parameters
......@@ -180,15 +205,16 @@ def _gen_copy_reduce_func(binary_op, reduce_op):
-----
This function supports autograd (computing input gradients given the output gradient).
""".format(
binary_str[binary_op],
binary_str[binary_op], reduce_op, x_str[binary_op]
),
reduce_op,
x_str[binary_op]), reduce_op)
)
def func(g, x):
if binary_op == 'copy_u':
return gspmm(g, 'copy_lhs', reduce_op, x, None)
if binary_op == "copy_u":
return gspmm(g, "copy_lhs", reduce_op, x, None)
else:
return gspmm(g, 'copy_rhs', reduce_op, None, x)
return gspmm(g, "copy_rhs", reduce_op, None, x)
func.__name__ = name
func.__doc__ = docstring(binary_op)
......
python/dgl/optim/__init__.py
View file @ 76bb5404
"""dgl optims."""
import importlib
import sys
import os
import sys
from ..backend import backend_name
from ..utils import expand_as_pair
def _load_backend(mod_name):
mod = importlib.import_module('.%s' % mod_name, __name__)
mod = importlib.import_module(".%s" % mod_name, __name__)
thismod = sys.modules[__name__]
for api, obj in mod.__dict__.items():
setattr(thismod, api, obj)
_load_backend(backend_name)
python/dgl/optim/pytorch/sparse_optim.py
View file @ 76bb5404
"""Node embedding optimizers"""
import abc
from abc import abstractmethod
import torch as th
from ...utils import get_shared_mem_array, create_shared_mem_array, \
pin_memory_inplace, gather_pinned_tensor_rows, \
scatter_pinned_tensor_rows
from ...nn.pytorch import NodeEmbedding
from ...cuda import nccl
from ...nn.pytorch import NodeEmbedding
from ...partition import NDArrayPartition
from ...utils import (
create_shared_mem_array,
gather_pinned_tensor_rows,
get_shared_mem_array,
pin_memory_inplace,
scatter_pinned_tensor_rows,
)
class SparseGradOptimizer(abc.ABC):
r''' The abstract sparse optimizer.
r"""The abstract sparse optimizer.
Note: dgl sparse optimizer only work with dgl.NodeEmbedding
......@@ -21,7 +27,8 @@ class SparseGradOptimizer(abc.ABC):
The list of NodeEmbeddings.
lr : float
The learning rate.
'''
"""
def __init__(self, params, lr):
self._params = params
self._lr = lr
......@@ -37,47 +44,54 @@ class SparseGradOptimizer(abc.ABC):
# otherwise it will crash the training
self.shmem_buffer_holder = []
assert len(params) > 0, 'Empty parameters'
assert len(params) > 0, "Empty parameters"
# if we are using shared memory for communication
for emb in params:
assert isinstance(emb, NodeEmbedding), \
'DGL SparseOptimizer only supports dgl.nn.NodeEmbedding'
assert isinstance(
emb, NodeEmbedding
), "DGL SparseOptimizer only supports dgl.nn.NodeEmbedding"
if self._rank is None:
self._rank = emb.rank
self._world_size = emb.world_size
else:
assert self._rank == emb.rank, \
'MultiGPU rank for each embedding should be same.'
assert self._world_size == emb.world_size, \
'MultiGPU world_size for each embedding should be same.'
assert (
self._rank == emb.rank
), "MultiGPU rank for each embedding should be same."
assert (
self._world_size == emb.world_size
), "MultiGPU world_size for each embedding should be same."
assert not self._rank is None
assert not self._world_size is None
self._nccl_root_id = 'SparseGradOptimizer.nccl_root_id'
self._nccl_root_id = "SparseGradOptimizer.nccl_root_id"
def step(self):
''' The step function.
"""The step function.
The step function is invoked at the end of every batch to update embeddings
'''
"""
# on the first step, check to see if the grads are on the GPU
if self._first_step:
for emb in self._params:
for _, data in emb._trace:
if data.grad.data.device.type == 'cuda':
if data.grad.data.device.type == "cuda":
# create a communicator
if self._device:
assert self._device == data.grad.device, \
"All gradients must be on the same device"
assert (
self._device == data.grad.device
), "All gradients must be on the same device"
else:
self._device = data.grad.device
else:
assert not self._device, \
"All gradients must be on the same device"
assert (
not self._device
), "All gradients must be on the same device"
# distributed backend use nccl
if self._device and \
(not th.distributed.is_initialized() or th.distributed.get_backend() == 'nccl'):
if self._device and (
not th.distributed.is_initialized()
or th.distributed.get_backend() == "nccl"
):
# device is only set if the grads are on a GPU
self._comm_setup()
else:
......@@ -91,16 +105,15 @@ class SparseGradOptimizer(abc.ABC):
self._shared_step()
def setup(self, params):
''' This is function where subclasses can perform any setup they need
to. It will be called during the first step, and communicators or
shared memory will have been setup before this call.
Parameters
----------
params : list of NodeEmbedding
The list of NodeEmbeddings.
'''
"""This is function where subclasses can perform any setup they need
to. It will be called during the first step, and communicators or
shared memory will have been setup before this call.
Parameters
----------
params : list of NodeEmbedding
The list of NodeEmbeddings.
"""
def _comm_setup(self):
# find a store to communicate the unique id through
......@@ -120,26 +133,32 @@ class SparseGradOptimizer(abc.ABC):
uid = store.get(self._nccl_root_id)
nccl_id = nccl.UniqueId(uid)
# needs to be set for nccl to work
self._comm = nccl.Communicator(self._world_size,
self._rank,
nccl_id)
self._comm = nccl.Communicator(
self._world_size, self._rank, nccl_id
)
th.distributed.barrier()
def _shared_setup(self):
for emb in self._params:
emb_name = emb.name
if self._rank == 0: # the master gpu process
opt_meta = create_shared_mem_array(emb_name+'_opt_meta', \
(self._world_size, self._world_size), th.int32).zero_()
if self._rank == 0: # the master gpu process
opt_meta = create_shared_mem_array(
emb_name + "_opt_meta",
(self._world_size, self._world_size),
th.int32,
).zero_()
if self._rank == 0:
emb.store.set(emb_name+'_opt_meta', emb_name)
emb.store.set(emb_name + "_opt_meta", emb_name)
self._opt_meta[emb_name] = opt_meta
elif self._rank > 0:
# receive
emb.store.wait([emb_name+'_opt_meta'])
opt_meta = get_shared_mem_array(emb_name+'_opt_meta', \
(self._world_size, self._world_size), th.int32)
emb.store.wait([emb_name + "_opt_meta"])
opt_meta = get_shared_mem_array(
emb_name + "_opt_meta",
(self._world_size, self._world_size),
th.int32,
)
self._opt_meta[emb_name] = opt_meta
def _comm_step(self):
......@@ -147,7 +166,7 @@ class SparseGradOptimizer(abc.ABC):
with th.no_grad():
idx_in = {}
grad_in = {}
for emb in self._params: # pylint: disable=too-many-nested-blocks
for emb in self._params: # pylint: disable=too-many-nested-blocks
emb_name = emb.name
partition = emb.partition
......@@ -156,14 +175,17 @@ class SparseGradOptimizer(abc.ABC):
partition = NDArrayPartition(
emb.num_embeddings,
self._world_size if self._world_size > 0 else 1,
mode='remainder')
mode="remainder",
)
# we need to combine gradients from multiple forward paths
if len(emb._trace) == 0:
idx = th.zeros((0,), dtype=th.long, device=self._device)
grad = th.zeros((0, emb.embedding_dim),
dtype=th.float32,
device=self._device)
grad = th.zeros(
(0, emb.embedding_dim),
dtype=th.float32,
device=self._device,
)
elif len(emb._trace) == 1:
# the special case where we can use the tensors as is
# without any memcpy's
......@@ -178,9 +200,10 @@ class SparseGradOptimizer(abc.ABC):
idx = th.cat(idx, dim=0)
grad = th.cat(grad, dim=0)
idx_in[emb_name], grad_in[emb_name] = \
comm.sparse_all_to_all_push(
idx, grad, partition=partition)
(
idx_in[emb_name],
grad_in[emb_name],
) = comm.sparse_all_to_all_push(idx, grad, partition=partition)
if emb.partition:
# if the embedding is partitioned, map back to indexes
# into the local tensor
......@@ -198,7 +221,6 @@ class SparseGradOptimizer(abc.ABC):
grad = grad_in[emb_name]
self.update(idx, grad, emb)
def _shared_step(self):
with th.no_grad():
# Frequently alloc and free shared memory to hold intermediate tensor is expensive
......@@ -206,7 +228,7 @@ class SparseGradOptimizer(abc.ABC):
shared_emb = {emb.name: ([], []) for emb in self._params}
# Go through all sparse embeddings
for emb in self._params: # pylint: disable=too-many-nested-blocks
for emb in self._params: # pylint: disable=too-many-nested-blocks
emb_name = emb.name
# we need to combine gradients from multiple forward paths
......@@ -222,10 +244,20 @@ class SparseGradOptimizer(abc.ABC):
# Note: we cannot skip the gradient exchange and update steps as other
# working processes may send gradient update requests corresponding
# to certain embedding to this process.
idx = th.cat(idx, dim=0) if len(idx) != 0 else \
th.zeros((0,), dtype=th.long, device=th.device('cpu'))
grad = th.cat(grad, dim=0) if len(grad) != 0 else \
th.zeros((0, emb.embedding_dim), dtype=th.float32, device=th.device('cpu'))
idx = (
th.cat(idx, dim=0)
if len(idx) != 0
else th.zeros((0,), dtype=th.long, device=th.device("cpu"))
)
grad = (
th.cat(grad, dim=0)
if len(grad) != 0
else th.zeros(
(0, emb.embedding_dim),
dtype=th.float32,
device=th.device("cpu"),
)
)
device = grad.device
idx_dtype = idx.dtype
......@@ -263,22 +295,40 @@ class SparseGradOptimizer(abc.ABC):
# currently nccl does not support Alltoallv operation
# we need to use CPU shared memory to share gradient
# across processes
idx_i = idx_i.to(th.device('cpu'))
grad_i = grad_i.to(th.device('cpu'))
idx_shmem_name = 'idx_{}_{}_{}'.format(emb_name, self._rank, i)
grad_shmem_name = 'grad_{}_{}_{}'.format(emb_name, self._rank, i)
idx_i = idx_i.to(th.device("cpu"))
grad_i = grad_i.to(th.device("cpu"))
idx_shmem_name = "idx_{}_{}_{}".format(
emb_name, self._rank, i
)
grad_shmem_name = "grad_{}_{}_{}".format(
emb_name, self._rank, i
)
# Create shared memory to hold temporary index and gradient tensor for
# cross-process send and recv.
if idx_shmem_name not in self._shared_cache[emb_name] or \
self._shared_cache[emb_name][idx_shmem_name].shape[0] \
< idx_i.shape[0]:
if idx_shmem_name in self._shared_cache[emb_name]:
if (
idx_shmem_name
not in self._shared_cache[emb_name]
or self._shared_cache[emb_name][
idx_shmem_name
].shape[0]
< idx_i.shape[0]
):
if (
idx_shmem_name
in self._shared_cache[emb_name]
):
self.shmem_buffer_holder.append(
self._shared_cache[emb_name][idx_shmem_name])
self._shared_cache[emb_name][
idx_shmem_name
]
)
self.shmem_buffer_holder.append(
self._shared_cache[emb_name][grad_shmem_name])
self._shared_cache[emb_name][
grad_shmem_name
]
)
# The total number of buffers is the number of NodeEmbeddings *
# world_size * (world_size - 1). The minimun buffer size is 128.
......@@ -287,22 +337,40 @@ class SparseGradOptimizer(abc.ABC):
# frequent shared memory allocation.
# The overall buffer cost will be smaller than three times
# the maximum memory requirement for sharing gradients.
buffer_size = 128 if idx_i.shape[0] < 128 else idx_i.shape[0] * 2
buffer_size = (
128
if idx_i.shape[0] < 128
else idx_i.shape[0] * 2
)
idx_shmem = create_shared_mem_array(
'{}_{}'.format(idx_shmem_name, buffer_size), \
(buffer_size,), idx_dtype)
"{}_{}".format(idx_shmem_name, buffer_size),
(buffer_size,),
idx_dtype,
)
grad_shmem = create_shared_mem_array(
'{}_{}'.format(grad_shmem_name, buffer_size), \
(buffer_size, grad_dim), grad_dtype)
self._shared_cache[emb_name][idx_shmem_name] = idx_shmem
self._shared_cache[emb_name][grad_shmem_name] = grad_shmem
"{}_{}".format(
grad_shmem_name, buffer_size
),
(buffer_size, grad_dim),
grad_dtype,
)
self._shared_cache[emb_name][
idx_shmem_name
] = idx_shmem
self._shared_cache[emb_name][
grad_shmem_name
] = grad_shmem
# Fill shared memory with temporal index tensor and gradient tensor
self._shared_cache[emb_name][idx_shmem_name][:idx_i.shape[0]] \
= idx_i
self._shared_cache[emb_name][grad_shmem_name][:idx_i.shape[0]] \
= grad_i
self._opt_meta[emb_name][self._rank][i] = idx_i.shape[0]
self._shared_cache[emb_name][idx_shmem_name][
: idx_i.shape[0]
] = idx_i
self._shared_cache[emb_name][grad_shmem_name][
: idx_i.shape[0]
] = grad_i
self._opt_meta[emb_name][self._rank][
i
] = idx_i.shape[0]
else:
shared_emb[emb_name][0].append(idx)
shared_emb[emb_name][1].append(grad)
......@@ -310,7 +378,7 @@ class SparseGradOptimizer(abc.ABC):
# make sure the idx shape is passed to each process through opt_meta
if self._world_size > 1:
th.distributed.barrier()
for emb in self._params: # pylint: disable=too-many-nested-blocks
for emb in self._params: # pylint: disable=too-many-nested-blocks
emb_name = emb.name
if self._world_size > 1:
# The first element in shared_emb[emb_name][0] is the local idx
......@@ -318,30 +386,56 @@ class SparseGradOptimizer(abc.ABC):
# gather gradients from all other processes
for i in range(self._world_size):
if i != self._rank:
idx_shmem_name = 'idx_{}_{}_{}'.format(emb_name, i, self._rank)
grad_shmem_name = 'grad_{}_{}_{}'.format(emb_name, i, self._rank)
idx_shmem_name = "idx_{}_{}_{}".format(
emb_name, i, self._rank
)
grad_shmem_name = "grad_{}_{}_{}".format(
emb_name, i, self._rank
)
size = self._opt_meta[emb_name][i][self._rank]
# Retrive shared memory holding the temporal index and gradient
# tensor that is sent to current training process
if idx_shmem_name not in self._shared_cache[emb_name] or \
self._shared_cache[emb_name][idx_shmem_name].shape[0] < size:
if (
idx_shmem_name
not in self._shared_cache[emb_name]
or self._shared_cache[emb_name][
idx_shmem_name
].shape[0]
< size
):
buffer_size = 128 if size < 128 else size * 2
idx_shmem = get_shared_mem_array(
'{}_{}'.format(idx_shmem_name, buffer_size), \
(buffer_size,), idx_dtype)
"{}_{}".format(idx_shmem_name, buffer_size),
(buffer_size,),
idx_dtype,
)
grad_shmem = get_shared_mem_array(
'{}_{}'.format(grad_shmem_name, buffer_size), \
(buffer_size, grad_dim), grad_dtype)
self._shared_cache[emb_name][idx_shmem_name] = idx_shmem
self._shared_cache[emb_name][grad_shmem_name] = grad_shmem
idx_i = self._shared_cache[emb_name][idx_shmem_name][:size]
grad_i = self._shared_cache[emb_name][grad_shmem_name][:size]
shared_emb[emb_name][0].append(idx_i.to(device,
non_blocking=True))
shared_emb[emb_name][1].append(grad_i.to(device,
non_blocking=True))
"{}_{}".format(
grad_shmem_name, buffer_size
),
(buffer_size, grad_dim),
grad_dtype,
)
self._shared_cache[emb_name][
idx_shmem_name
] = idx_shmem
self._shared_cache[emb_name][
grad_shmem_name
] = grad_shmem
idx_i = self._shared_cache[emb_name][
idx_shmem_name
][:size]
grad_i = self._shared_cache[emb_name][
grad_shmem_name
][:size]
shared_emb[emb_name][0].append(
idx_i.to(device, non_blocking=True)
)
shared_emb[emb_name][1].append(
grad_i.to(device, non_blocking=True)
)
if self._clean_grad:
# clean gradient track
......@@ -362,7 +456,7 @@ class SparseGradOptimizer(abc.ABC):
@abstractmethod
def update(self, idx, grad, emb):
""" Update embeddings in a sparse manner
"""Update embeddings in a sparse manner
Sparse embeddings are updated in mini batches. We maintain gradient states for
each embedding so they can be updated separately.
......@@ -377,12 +471,12 @@ class SparseGradOptimizer(abc.ABC):
"""
def zero_grad(self):
"""clean grad cache
"""
"""clean grad cache"""
self._clean_grad = True
class SparseAdagrad(SparseGradOptimizer):
r''' Node embedding optimizer using the Adagrad algorithm.
r"""Node embedding optimizer using the Adagrad algorithm.
This optimizer implements a sparse version of Adagrad algorithm for
optimizing :class:`dgl.nn.NodeEmbedding`. Being sparse means it only updates
......@@ -418,7 +512,8 @@ class SparseAdagrad(SparseGradOptimizer):
... loss = F.sum(feats + 1, 0)
... loss.backward()
... optimizer.step()
'''
"""
def __init__(self, params, lr, eps=1e-10):
super(SparseAdagrad, self).__init__(params, lr)
self._eps = eps
......@@ -426,38 +521,43 @@ class SparseAdagrad(SparseGradOptimizer):
def setup(self, params):
# We need to register a state sum for each embedding in the kvstore.
for emb in params:
assert isinstance(emb, NodeEmbedding), \
'SparseAdagrad only supports dgl.nn.NodeEmbedding'
assert isinstance(
emb, NodeEmbedding
), "SparseAdagrad only supports dgl.nn.NodeEmbedding"
emb_name = emb.name
if th.device(emb.emb_tensor.device) == th.device('cpu'):
if th.device(emb.emb_tensor.device) == th.device("cpu"):
# if our embedding is on the CPU, our state also has to be
if self._rank < 0:
state = th.empty(
emb.weight.shape,
dtype=th.float32,
device=th.device('cpu')).zero_()
device=th.device("cpu"),
).zero_()
elif self._rank == 0:
state = create_shared_mem_array(emb_name+'_state', \
emb.weight.shape, th.float32).zero_()
state = create_shared_mem_array(
emb_name + "_state", emb.weight.shape, th.float32
).zero_()
if self._world_size > 1:
emb.store.set(emb_name+'_opt', emb_name)
emb.store.set(emb_name + "_opt", emb_name)
elif self._rank > 0:
# receive
emb.store.wait([emb_name+'_opt'])
state = get_shared_mem_array(emb_name+'_state', \
emb.weight.shape, th.float32)
emb.store.wait([emb_name + "_opt"])
state = get_shared_mem_array(
emb_name + "_state", emb.weight.shape, th.float32
)
else:
# distributed state on on gpu
state = th.empty(
emb.emb_tensor.shape,
dtype=th.float32,
device=emb.emb_tensor.device).zero_()
device=emb.emb_tensor.device,
).zero_()
emb.set_optm_state(state)
def update(self, idx, grad, emb):
""" Update embeddings in a sparse manner
"""Update embeddings in a sparse manner
Sparse embeddings are updated in mini batches. We maintain gradient states for
each embedding so they can be updated separately.
......@@ -474,12 +574,16 @@ class SparseAdagrad(SparseGradOptimizer):
clr = self._lr
# the update is non-linear so indices must be unique
grad_indices, inverse, cnt = th.unique(idx, return_inverse=True, return_counts=True)
grad_values = th.zeros((grad_indices.shape[0], grad.shape[1]), device=grad.device)
grad_indices, inverse, cnt = th.unique(
idx, return_inverse=True, return_counts=True
)
grad_values = th.zeros(
(grad_indices.shape[0], grad.shape[1]), device=grad.device
)
grad_values.index_add_(0, inverse, grad)
grad_values = grad_values / cnt.unsqueeze(1)
grad_sum = (grad_values * grad_values)
grad_sum = grad_values * grad_values
state = emb.optm_state
state_dev = state.device
state_idx = grad_indices.to(state_dev)
......@@ -491,8 +595,9 @@ class SparseAdagrad(SparseGradOptimizer):
tmp = clr * grad_values / std_values
emb.weight[state_idx] -= tmp.to(state_dev)
class SparseAdam(SparseGradOptimizer):
r''' Node embedding optimizer using the Adam algorithm.
r"""Node embedding optimizer using the Adam algorithm.
This optimizer implements a sparse version of Adagrad algorithm for
optimizing :class:`dgl.nn.NodeEmbedding`. Being sparse means it only
......@@ -545,9 +650,17 @@ class SparseAdam(SparseGradOptimizer):
... loss = F.sum(feats + 1, 0)
... loss.backward()
... optimizer.step()
'''
def __init__(self, params, lr, betas=(0.9, 0.999), eps=1e-08, \
use_uva=None, dtype=th.float32):
"""
def __init__(
self,
params,
lr,
betas=(0.9, 0.999),
eps=1e-08,
use_uva=None,
dtype=th.float32,
):
super(SparseAdam, self).__init__(params, lr)
self._lr = lr
self._beta1 = betas[0]
......@@ -556,9 +669,10 @@ class SparseAdam(SparseGradOptimizer):
self._use_uva = use_uva
self._nd_handle = {}
self._is_using_uva = {}
assert dtype in [th.float16, th.float32], \
"Unsupported dtype {}. Valid choices are th.float32 " \
assert dtype in [th.float16, th.float32], (
"Unsupported dtype {}. Valid choices are th.float32 "
"and th.float32".format(dtype)
)
self._dtype = dtype
def _setup_uva(self, name, mem, power):
......@@ -570,44 +684,54 @@ class SparseAdam(SparseGradOptimizer):
def setup(self, params):
# We need to register a state sum for each embedding in the kvstore.
for emb in params:
assert isinstance(emb, NodeEmbedding), \
'SparseAdam only supports dgl.nn.NodeEmbedding'
assert isinstance(
emb, NodeEmbedding
), "SparseAdam only supports dgl.nn.NodeEmbedding"
emb_name = emb.name
self._is_using_uva[emb_name] = self._use_uva
if th.device(emb.emb_tensor.device) == th.device('cpu'):
if th.device(emb.emb_tensor.device) == th.device("cpu"):
# if our embedding is on the CPU, our state also has to be
if self._rank < 0:
state_step = th.empty(
(emb.weight.shape[0],),
dtype=th.int32,
device=th.device('cpu')).zero_()
device=th.device("cpu"),
).zero_()
state_mem = th.empty(
emb.weight.shape,
dtype=self._dtype,
device=th.device('cpu')).zero_()
device=th.device("cpu"),
).zero_()
state_power = th.empty(
emb.weight.shape,
dtype=self._dtype,
device=th.device('cpu')).zero_()
device=th.device("cpu"),
).zero_()
elif self._rank == 0:
state_step = create_shared_mem_array(emb_name+'_step', \
(emb.weight.shape[0],), th.int32).zero_()
state_mem = create_shared_mem_array(emb_name+'_mem', \
emb.weight.shape, self._dtype).zero_()
state_power = create_shared_mem_array(emb_name+'_power', \
emb.weight.shape, self._dtype).zero_()
state_step = create_shared_mem_array(
emb_name + "_step", (emb.weight.shape[0],), th.int32
).zero_()
state_mem = create_shared_mem_array(
emb_name + "_mem", emb.weight.shape, self._dtype
).zero_()
state_power = create_shared_mem_array(
emb_name + "_power", emb.weight.shape, self._dtype
).zero_()
if self._world_size > 1:
emb.store.set(emb_name+'_opt', emb_name)
emb.store.set(emb_name + "_opt", emb_name)
elif self._rank > 0:
# receive
emb.store.wait([emb_name+'_opt'])
state_step = get_shared_mem_array(emb_name+'_step', \
(emb.weight.shape[0],), th.int32)
state_mem = get_shared_mem_array(emb_name+'_mem', \
emb.weight.shape, self._dtype)
state_power = get_shared_mem_array(emb_name+'_power', \
emb.weight.shape, self._dtype)
emb.store.wait([emb_name + "_opt"])
state_step = get_shared_mem_array(
emb_name + "_step", (emb.weight.shape[0],), th.int32
)
state_mem = get_shared_mem_array(
emb_name + "_mem", emb.weight.shape, self._dtype
)
state_power = get_shared_mem_array(
emb_name + "_power", emb.weight.shape, self._dtype
)
if self._is_using_uva[emb_name]:
# if use_uva has been explicitly set to true, otherwise
......@@ -621,20 +745,23 @@ class SparseAdam(SparseGradOptimizer):
state_step = th.empty(
[emb.emb_tensor.shape[0]],
dtype=th.int32,
device=emb.emb_tensor.device).zero_()
device=emb.emb_tensor.device,
).zero_()
state_mem = th.empty(
emb.emb_tensor.shape,
dtype=self._dtype,
device=emb.emb_tensor.device).zero_()
device=emb.emb_tensor.device,
).zero_()
state_power = th.empty(
emb.emb_tensor.shape,
dtype=self._dtype,
device=emb.emb_tensor.device).zero_()
device=emb.emb_tensor.device,
).zero_()
state = (state_step, state_mem, state_power)
emb.set_optm_state(state)
def update(self, idx, grad, emb):
""" Update embeddings in a sparse manner
"""Update embeddings in a sparse manner
Sparse embeddings are updated in mini batches. We maintain gradient states for
each embedding so they can be updated separately.
......@@ -655,7 +782,7 @@ class SparseAdam(SparseGradOptimizer):
# whether or not we need to transfer data from the GPU to the CPU
# while updating the weights
is_d2h = state_dev.type == 'cpu' and exec_dev.type == 'cuda'
is_d2h = state_dev.type == "cpu" and exec_dev.type == "cuda"
# only perform async copies cpu -> gpu, or gpu-> gpu, but block
# when copying to the cpu, so as to ensure the copy is finished
......@@ -678,16 +805,18 @@ class SparseAdam(SparseGradOptimizer):
clr = self._lr
# There can be duplicated indices due to sampling.
# Thus unique them here and average the gradient here.
grad_indices, inverse, cnt = th.unique(idx,
return_inverse=True,
return_counts=True)
grad_indices, inverse, cnt = th.unique(
idx, return_inverse=True, return_counts=True
)
state_idx = grad_indices.to(state_dev)
state_step[state_idx] += 1
state_step = state_step[state_idx].to(exec_dev)
if use_uva:
orig_mem = gather_pinned_tensor_rows(state_mem, grad_indices)
orig_power = gather_pinned_tensor_rows(state_power, grad_indices)
orig_power = gather_pinned_tensor_rows(
state_power, grad_indices
)
else:
orig_mem = state_mem[state_idx].to(exec_dev)
orig_power = state_power[state_idx].to(exec_dev)
......@@ -695,38 +824,48 @@ class SparseAdam(SparseGradOptimizer):
orig_mem = orig_mem.to(dtype=exec_dtype)
orig_power = orig_power.to(dtype=exec_dtype)
grad_values = th.zeros((grad_indices.shape[0], grad.shape[1]), device=exec_dev)
grad_values = th.zeros(
(grad_indices.shape[0], grad.shape[1]), device=exec_dev
)
grad_values.index_add_(0, inverse, grad)
grad_values = grad_values / cnt.unsqueeze(1)
grad_mem = grad_values
grad_power = grad_values * grad_values
update_mem = beta1 * orig_mem + (1.-beta1) * grad_mem
update_power = beta2 * orig_power + (1.-beta2) * grad_power
update_mem = beta1 * orig_mem + (1.0 - beta1) * grad_mem
update_power = beta2 * orig_power + (1.0 - beta2) * grad_power
if use_uva:
scatter_pinned_tensor_rows(state_mem, \
grad_indices, \
update_mem.to(dtype=self._dtype))
scatter_pinned_tensor_rows(state_power, \
grad_indices, \
update_power.to(dtype=self._dtype))
scatter_pinned_tensor_rows(
state_mem, grad_indices, update_mem.to(dtype=self._dtype)
)
scatter_pinned_tensor_rows(
state_power,
grad_indices,
update_power.to(dtype=self._dtype),
)
else:
update_mem_dst = update_mem.to(dtype=self._dtype).to(
state_dev, non_blocking=True)
state_dev, non_blocking=True
)
update_power_dst = update_power.to(dtype=self._dtype).to(
state_dev, non_blocking=True)
state_dev, non_blocking=True
)
if state_block:
# use events to try and overlap CPU and GPU as much as possible
update_event = th.cuda.Event()
update_event.record()
update_mem_corr = update_mem / (1. - th.pow(th.tensor(beta1, device=exec_dev),
state_step)).unsqueeze(1)
update_power_corr = update_power / (1. - th.pow(th.tensor(beta2, device=exec_dev),
state_step)).unsqueeze(1)
std_values = clr * update_mem_corr / (th.sqrt(update_power_corr) + eps)
update_mem_corr = update_mem / (
1.0 - th.pow(th.tensor(beta1, device=exec_dev), state_step)
).unsqueeze(1)
update_power_corr = update_power / (
1.0 - th.pow(th.tensor(beta2, device=exec_dev), state_step)
).unsqueeze(1)
std_values = (
clr * update_mem_corr / (th.sqrt(update_power_corr) + eps)
)
std_values_dst = std_values.to(state_dev, non_blocking=True)
if state_block:
......
python/dgl/partition.py
View file @ 76bb5404
......@@ -2,22 +2,26 @@
import os
import re
import time
import numpy as np
from ._ffi.function import _init_api
from .heterograph import DGLHeteroGraph
from . import backend as F
from . import utils
from ._ffi.function import _init_api
from .base import EID, ETYPE, NID, NTYPE
from .heterograph import DGLHeteroGraph
from .ndarray import NDArray
from .base import EID, NID, NTYPE, ETYPE
from .subgraph import edge_subgraph
__all__ = ["metis_partition", "metis_partition_assignment",
"partition_graph_with_halo"]
__all__ = [
"metis_partition",
"metis_partition_assignment",
"partition_graph_with_halo",
]
def reorder_nodes(g, new_node_ids):
""" Generate a new graph with new node IDs.
"""Generate a new graph with new node IDs.
We assign each node in the input graph with a new node ID. This results in
a new graph.
......@@ -33,25 +37,29 @@ def reorder_nodes(g, new_node_ids):
DGLGraph
The graph with new node IDs.
"""
assert len(new_node_ids) == g.number_of_nodes(), \
"The number of new node ids must match #nodes in the graph."
assert (
len(new_node_ids) == g.number_of_nodes()
), "The number of new node ids must match #nodes in the graph."
new_node_ids = utils.toindex(new_node_ids)
sorted_ids, idx = F.sort_1d(new_node_ids.tousertensor())
assert F.asnumpy(sorted_ids[0]) == 0 \
and F.asnumpy(sorted_ids[-1]) == g.number_of_nodes() - 1, \
"The new node IDs are incorrect."
assert (
F.asnumpy(sorted_ids[0]) == 0
and F.asnumpy(sorted_ids[-1]) == g.number_of_nodes() - 1
), "The new node IDs are incorrect."
new_gidx = _CAPI_DGLReorderGraph_Hetero(
g._graph, new_node_ids.todgltensor())
new_g = DGLHeteroGraph(gidx=new_gidx, ntypes=['_N'], etypes=['_E'])
new_g.ndata['orig_id'] = idx
g._graph, new_node_ids.todgltensor()
)
new_g = DGLHeteroGraph(gidx=new_gidx, ntypes=["_N"], etypes=["_E"])
new_g.ndata["orig_id"] = idx
return new_g
def _get_halo_heterosubgraph_inner_node(halo_subg):
return _CAPI_GetHaloSubgraphInnerNodes_Hetero(halo_subg)
def reshuffle_graph(g, node_part=None):
'''Reshuffle node ids and edge IDs of a graph.
"""Reshuffle node ids and edge IDs of a graph.
This function reshuffles nodes and edges in a graph so that all nodes/edges of the same type
have contiguous IDs. If a graph is partitioned and nodes are assigned to different partitions,
......@@ -71,11 +79,11 @@ def reshuffle_graph(g, node_part=None):
(DGLGraph, Tensor)
The graph whose nodes and edges are reshuffled.
The 1D tensor that indicates the partition IDs of the nodes in the reshuffled graph.
'''
"""
# In this case, we don't need to reshuffle node IDs and edge IDs.
if node_part is None:
g.ndata['orig_id'] = F.arange(0, g.number_of_nodes())
g.edata['orig_id'] = F.arange(0, g.number_of_edges())
g.ndata["orig_id"] = F.arange(0, g.number_of_nodes())
g.edata["orig_id"] = F.arange(0, g.number_of_edges())
return g, None
start = time.time()
......@@ -89,38 +97,48 @@ def reshuffle_graph(g, node_part=None):
if is_hetero:
num_node_types = F.max(g.ndata[NTYPE], 0) + 1
if node_part is not None:
sorted_part, new2old_map = F.sort_1d(node_part * num_node_types + g.ndata[NTYPE])
sorted_part, new2old_map = F.sort_1d(
node_part * num_node_types + g.ndata[NTYPE]
)
else:
sorted_part, new2old_map = F.sort_1d(g.ndata[NTYPE])
sorted_part = F.floor_div(sorted_part, num_node_types)
elif node_part is not None:
sorted_part, new2old_map = F.sort_1d(node_part)
else:
g.ndata['orig_id'] = g.ndata[NID]
g.edata['orig_id'] = g.edata[EID]
g.ndata["orig_id"] = g.ndata[NID]
g.edata["orig_id"] = g.edata[EID]
return g, None
new_node_ids = np.zeros((g.number_of_nodes(),), dtype=np.int64)
new_node_ids[F.asnumpy(new2old_map)] = np.arange(0, g.number_of_nodes())
# If the input graph is homogneous, we only need to create an empty array, so that
# _CAPI_DGLReassignEdges_Hetero knows how to handle it.
etype = g.edata[ETYPE] if ETYPE in g.edata else F.zeros((0), F.dtype(sorted_part), F.cpu())
etype = (
g.edata[ETYPE]
if ETYPE in g.edata
else F.zeros((0), F.dtype(sorted_part), F.cpu())
)
g = reorder_nodes(g, new_node_ids)
node_part = utils.toindex(sorted_part)
# We reassign edges in in-CSR. In this way, after partitioning, we can ensure
# that all edges in a partition are in the contiguous ID space.
etype_idx = utils.toindex(etype)
orig_eids = _CAPI_DGLReassignEdges_Hetero(g._graph, etype_idx.todgltensor(),
node_part.todgltensor(), True)
orig_eids = _CAPI_DGLReassignEdges_Hetero(
g._graph, etype_idx.todgltensor(), node_part.todgltensor(), True
)
orig_eids = utils.toindex(orig_eids)
orig_eids = orig_eids.tousertensor()
g.edata['orig_id'] = orig_eids
g.edata["orig_id"] = orig_eids
print('Reshuffle nodes and edges: {:.3f} seconds'.format(time.time() - start))
print(
"Reshuffle nodes and edges: {:.3f} seconds".format(time.time() - start)
)
return g, node_part.tousertensor()
def partition_graph_with_halo(g, node_part, extra_cached_hops, reshuffle=False):
'''Partition a graph.
"""Partition a graph.
Based on the given node assignments for each partition, the function splits
the input graph into subgraphs. A subgraph may contain HALO nodes which does
......@@ -156,20 +174,21 @@ def partition_graph_with_halo(g, node_part, extra_cached_hops, reshuffle=False):
Tensor
1D tensor that stores the mapping between the reshuffled edge IDs and
the original edge IDs if 'reshuffle=True'. Otherwise, return None.
'''
"""
assert len(node_part) == g.number_of_nodes()
if reshuffle:
g, node_part = reshuffle_graph(g, node_part)
orig_nids = g.ndata['orig_id']
orig_eids = g.edata['orig_id']
orig_nids = g.ndata["orig_id"]
orig_eids = g.edata["orig_id"]
node_part = utils.toindex(node_part)
start = time.time()
subgs = _CAPI_DGLPartitionWithHalo_Hetero(
g._graph, node_part.todgltensor(), extra_cached_hops)
g._graph, node_part.todgltensor(), extra_cached_hops
)
# g is no longer needed. Free memory.
g = None
print('Split the graph: {:.3f} seconds'.format(time.time() - start))
print("Split the graph: {:.3f} seconds".format(time.time() - start))
subg_dict = {}
node_part = node_part.tousertensor()
start = time.time()
......@@ -182,14 +201,17 @@ def partition_graph_with_halo(g, node_part, extra_cached_hops, reshuffle=False):
inner_nids = F.nonzero_1d(inner_node)
# TODO(zhengda) we need to fix utils.toindex() to avoid the dtype cast below.
inner_nids = F.astype(inner_nids, F.int64)
inner_eids = subg.in_edges(inner_nids, form='eid')
inner_edge = F.scatter_row(inner_edge, inner_eids,
F.ones((len(inner_eids),), F.dtype(inner_edge), F.cpu()))
inner_eids = subg.in_edges(inner_nids, form="eid")
inner_edge = F.scatter_row(
inner_edge,
inner_eids,
F.ones((len(inner_eids),), F.dtype(inner_edge), F.cpu()),
)
return inner_edge
# This creaets a subgraph from subgraphs returned from the CAPI above.
def create_subgraph(subg, induced_nodes, induced_edges, inner_node):
subg1 = DGLHeteroGraph(gidx=subg.graph, ntypes=['_N'], etypes=['_E'])
subg1 = DGLHeteroGraph(gidx=subg.graph, ntypes=["_N"], etypes=["_E"])
# If IDs are shuffled, we should shuffled edges. This will help us collect edge data
# from the distributed graph after training.
if reshuffle:
......@@ -199,7 +221,9 @@ def partition_graph_with_halo(g, node_part, extra_cached_hops, reshuffle=False):
# of outer edges with a large value, so we will get the sorted list as we want.
max_eid = F.max(induced_edges[0], 0) + 1
inner_edge = get_inner_edge(subg1, inner_node)
eid = F.astype(induced_edges[0], F.int64) + max_eid * F.astype(inner_edge == 0, F.int64)
eid = F.astype(induced_edges[0], F.int64) + max_eid * F.astype(
inner_edge == 0, F.int64
)
_, index = F.sort_1d(eid)
subg1 = edge_subgraph(subg1, index, relabel_nodes=False)
......@@ -213,44 +237,49 @@ def partition_graph_with_halo(g, node_part, extra_cached_hops, reshuffle=False):
for i, subg in enumerate(subgs):
inner_node = _get_halo_heterosubgraph_inner_node(subg)
inner_node = F.zerocopy_from_dlpack(inner_node.to_dlpack())
subg = create_subgraph(subg, subg.induced_nodes, subg.induced_edges, inner_node)
subg.ndata['inner_node'] = inner_node
subg.ndata['part_id'] = F.gather_row(node_part, subg.ndata[NID])
subg = create_subgraph(
subg, subg.induced_nodes, subg.induced_edges, inner_node
)
subg.ndata["inner_node"] = inner_node
subg.ndata["part_id"] = F.gather_row(node_part, subg.ndata[NID])
if reshuffle:
subg.ndata['orig_id'] = F.gather_row(orig_nids, subg.ndata[NID])
subg.edata['orig_id'] = F.gather_row(orig_eids, subg.edata[EID])
subg.ndata["orig_id"] = F.gather_row(orig_nids, subg.ndata[NID])
subg.edata["orig_id"] = F.gather_row(orig_eids, subg.edata[EID])
if extra_cached_hops >= 1:
inner_edge = get_inner_edge(subg, inner_node)
else:
inner_edge = F.ones((subg.number_of_edges(),), F.int8, F.cpu())
subg.edata['inner_edge'] = inner_edge
subg.edata["inner_edge"] = inner_edge
subg_dict[i] = subg
print('Construct subgraphs: {:.3f} seconds'.format(time.time() - start))
print("Construct subgraphs: {:.3f} seconds".format(time.time() - start))
if reshuffle:
return subg_dict, orig_nids, orig_eids
else:
return subg_dict, None, None
def get_peak_mem():
''' Get the peak memory size.
"""Get the peak memory size.
Returns
-------
float
The peak memory size in GB.
'''
if not os.path.exists('/proc/self/status'):
"""
if not os.path.exists("/proc/self/status"):
return 0.0
for line in open('/proc/self/status', 'r'):
if 'VmPeak' in line:
mem = re.findall(r'\d+', line)[0]
for line in open("/proc/self/status", "r"):
if "VmPeak" in line:
mem = re.findall(r"\d+", line)[0]
return int(mem) / 1024 / 1024
return 0.0
def metis_partition_assignment(g, k, balance_ntypes=None, balance_edges=False,
mode="k-way", objtype='cut'):
''' This assigns nodes to different partitions with Metis partitioning algorithm.
def metis_partition_assignment(
g, k, balance_ntypes=None, balance_edges=False, mode="k-way", objtype="cut"
):
"""This assigns nodes to different partitions with Metis partitioning algorithm.
When performing Metis partitioning, we can put some constraint on the partitioning.
Current, it supports two constrants to balance the partitioning. By default, Metis
......@@ -284,16 +313,24 @@ def metis_partition_assignment(g, k, balance_ntypes=None, balance_edges=False,
-------
a 1-D tensor
A vector with each element that indicates the partition ID of a vertex.
'''
assert mode in ("k-way", "recursive"), "'mode' can only be 'k-way' or 'recursive'"
assert g.idtype == F.int64, "IdType of graph is required to be int64 for now."
"""
assert mode in (
"k-way",
"recursive",
), "'mode' can only be 'k-way' or 'recursive'"
assert (
g.idtype == F.int64
), "IdType of graph is required to be int64 for now."
# METIS works only on symmetric graphs.
# The METIS runs on the symmetric graph to generate the node assignment to partitions.
start = time.time()
sym_gidx = _CAPI_DGLMakeSymmetric_Hetero(g._graph)
sym_g = DGLHeteroGraph(gidx=sym_gidx)
print('Convert a graph into a bidirected graph: {:.3f} seconds, peak memory: {:.3f} GB'.format(
time.time() - start, get_peak_mem()))
print(
"Convert a graph into a bidirected graph: {:.3f} seconds, peak memory: {:.3f} GB".format(
time.time() - start, get_peak_mem()
)
)
vwgt = []
# To balance the node types in each partition, we can take advantage of the vertex weights
# in Metis. When vertex weights are provided, Metis will tries to generate partitions with
......@@ -307,8 +344,9 @@ def metis_partition_assignment(g, k, balance_ntypes=None, balance_edges=False,
# of weights is the same as the number of node types.
start = time.time()
if balance_ntypes is not None:
assert len(balance_ntypes) == g.number_of_nodes(), \
"The length of balance_ntypes should be equal to #nodes in the graph"
assert (
len(balance_ntypes) == g.number_of_nodes()
), "The length of balance_ntypes should be equal to #nodes in the graph"
balance_ntypes = F.tensor(balance_ntypes)
uniq_ntypes = F.unique(balance_ntypes)
for ntype in uniq_ntypes:
......@@ -328,19 +366,31 @@ def metis_partition_assignment(g, k, balance_ntypes=None, balance_edges=False,
# The vertex weights have to be stored in a vector.
if len(vwgt) > 0:
vwgt = F.stack(vwgt, 1)
shape = (np.prod(F.shape(vwgt),),)
shape = (
np.prod(
F.shape(vwgt),
),
)
vwgt = F.reshape(vwgt, shape)
vwgt = F.to_dgl_nd(vwgt)
else:
vwgt = F.zeros((0,), F.int64, F.cpu())
vwgt = F.to_dgl_nd(vwgt)
print('Construct multi-constraint weights: {:.3f} seconds, peak memory: {:.3f} GB'.format(
time.time() - start, get_peak_mem()))
print(
"Construct multi-constraint weights: {:.3f} seconds, peak memory: {:.3f} GB".format(
time.time() - start, get_peak_mem()
)
)
start = time.time()
node_part = _CAPI_DGLMetisPartition_Hetero(sym_g._graph, k, vwgt, mode, (objtype == 'cut'))
print('Metis partitioning: {:.3f} seconds, peak memory: {:.3f} GB'.format(
time.time() - start, get_peak_mem()))
node_part = _CAPI_DGLMetisPartition_Hetero(
sym_g._graph, k, vwgt, mode, (objtype == "cut")
)
print(
"Metis partitioning: {:.3f} seconds, peak memory: {:.3f} GB".format(
time.time() - start, get_peak_mem()
)
)
if len(node_part) == 0:
return None
else:
......@@ -348,9 +398,16 @@ def metis_partition_assignment(g, k, balance_ntypes=None, balance_edges=False,
return node_part.tousertensor()
def metis_partition(g, k, extra_cached_hops=0, reshuffle=False,
balance_ntypes=None, balance_edges=False, mode="k-way"):
''' This is to partition a graph with Metis partitioning.
def metis_partition(
g,
k,
extra_cached_hops=0,
reshuffle=False,
balance_ntypes=None,
balance_edges=False,
mode="k-way",
):
"""This is to partition a graph with Metis partitioning.
Metis assigns vertices to partitions. This API constructs subgraphs with the vertices assigned
to the partitions and their incoming edges. A subgraph may contain HALO nodes which does
......@@ -398,18 +455,25 @@ def metis_partition(g, k, extra_cached_hops=0, reshuffle=False,
--------
a dict of DGLGraphs
The key is the partition ID and the value is the DGLGraph of the partition.
'''
assert mode in ("k-way", "recursive"), "'mode' can only be 'k-way' or 'recursive'"
node_part = metis_partition_assignment(g, k, balance_ntypes, balance_edges, mode)
"""
assert mode in (
"k-way",
"recursive",
), "'mode' can only be 'k-way' or 'recursive'"
node_part = metis_partition_assignment(
g, k, balance_ntypes, balance_edges, mode
)
if node_part is None:
return None
# Then we split the original graph into parts based on the METIS partitioning results.
return partition_graph_with_halo(g, node_part, extra_cached_hops, reshuffle)[0]
return partition_graph_with_halo(
g, node_part, extra_cached_hops, reshuffle
)[0]
class NDArrayPartition(object):
""" Create a new partition of an NDArray. That is, an object which assigns
"""Create a new partition of an NDArray. That is, an object which assigns
each row of an NDArray to a specific partition.
Parameters
......@@ -456,69 +520,78 @@ class NDArrayPartition(object):
>>> part = NDArrayPartition(g.num_nodes(), num_parts, mode='range',
... part_ranges=part_range)
"""
def __init__(self, array_size, num_parts, mode='remainder', part_ranges=None):
def __init__(
self, array_size, num_parts, mode="remainder", part_ranges=None
):
assert num_parts > 0, 'Invalid "num_parts", must be > 0.'
if mode == 'remainder':
assert part_ranges is None, 'When using remainder-based ' \
if mode == "remainder":
assert part_ranges is None, (
"When using remainder-based "
'partitioning, "part_ranges" should not be specified.'
)
self._partition = _CAPI_DGLNDArrayPartitionCreateRemainderBased(
array_size, num_parts)
elif mode == 'range':
assert part_ranges is not None, 'When using range-based ' \
array_size, num_parts
)
elif mode == "range":
assert part_ranges is not None, (
"When using range-based "
'partitioning, "part_ranges" must not be None.'
assert part_ranges[0] == 0 and part_ranges[-1] == array_size, \
'part_ranges[0] must be 0, and part_ranges[-1] must be ' \
)
assert part_ranges[0] == 0 and part_ranges[-1] == array_size, (
"part_ranges[0] must be 0, and part_ranges[-1] must be "
'"array_size".'
)
if F.is_tensor(part_ranges):
part_ranges = F.zerocopy_to_dgl_ndarray(part_ranges)
assert isinstance(part_ranges, NDArray), '"part_ranges" must ' \
'be Tensor or dgl.NDArray.'
assert isinstance(part_ranges, NDArray), (
'"part_ranges" must ' "be Tensor or dgl.NDArray."
)
self._partition = _CAPI_DGLNDArrayPartitionCreateRangeBased(
array_size,
num_parts,
part_ranges)
array_size, num_parts, part_ranges
)
else:
assert False, 'Unknown partition mode "{}"'.format(mode)
self._array_size = array_size
self._num_parts = num_parts
def num_parts(self):
""" Get the number of partitions.
"""
"""Get the number of partitions."""
return self._num_parts
def array_size(self):
""" Get the total size of the first dimension of the partitioned array.
"""
"""Get the total size of the first dimension of the partitioned array."""
return self._array_size
def get(self):
""" Get the C-handle for this object.
"""
"""Get the C-handle for this object."""
return self._partition
def get_local_indices(self, part, ctx):
""" Get the set of global indices in this given partition.
"""
return self.map_to_global(F.arange(0, self.local_size(part), ctx=ctx), part)
"""Get the set of global indices in this given partition."""
return self.map_to_global(
F.arange(0, self.local_size(part), ctx=ctx), part
)
def local_size(self, part):
""" Get the number of rows/items assigned to the given part.
"""
"""Get the number of rows/items assigned to the given part."""
return _CAPI_DGLNDArrayPartitionGetPartSize(self._partition, part)
def map_to_local(self, idxs):
""" Convert the set of global indices to local indices
"""
return F.zerocopy_from_dgl_ndarray(_CAPI_DGLNDArrayPartitionMapToLocal(
self._partition,
F.zerocopy_to_dgl_ndarray(idxs)))
"""Convert the set of global indices to local indices"""
return F.zerocopy_from_dgl_ndarray(
_CAPI_DGLNDArrayPartitionMapToLocal(
self._partition, F.zerocopy_to_dgl_ndarray(idxs)
)
)
def map_to_global(self, idxs, part_id):
""" Convert the set of local indices ot global indices
"""
return F.zerocopy_from_dgl_ndarray(_CAPI_DGLNDArrayPartitionMapToGlobal(
self._partition, F.zerocopy_to_dgl_ndarray(idxs), part_id))
"""Convert the set of local indices ot global indices"""
return F.zerocopy_from_dgl_ndarray(
_CAPI_DGLNDArrayPartitionMapToGlobal(
self._partition, F.zerocopy_to_dgl_ndarray(idxs), part_id
)
)
_init_api("dgl.partition")
python/dgl/propagate.py
View file @ 76bb5404
......@@ -5,14 +5,22 @@ from . import backend as F
from . import traversal as trv
from .heterograph import DGLHeteroGraph
__all__ = ['prop_nodes', 'prop_nodes_bfs', 'prop_nodes_topo',
'prop_edges', 'prop_edges_dfs']
def prop_nodes(graph,
nodes_generator,
message_func='default',
reduce_func='default',
apply_node_func='default'):
__all__ = [
"prop_nodes",
"prop_nodes_bfs",
"prop_nodes_topo",
"prop_edges",
"prop_edges_dfs",
]
def prop_nodes(
graph,
nodes_generator,
message_func="default",
reduce_func="default",
apply_node_func="default",
):
"""Functional method for :func:`dgl.DGLGraph.prop_nodes`.
Parameters
......@@ -30,13 +38,18 @@ def prop_nodes(graph,
--------
dgl.DGLGraph.prop_nodes
"""
graph.prop_nodes(nodes_generator, message_func, reduce_func, apply_node_func)
def prop_edges(graph,
edges_generator,
message_func='default',
reduce_func='default',
apply_node_func='default'):
graph.prop_nodes(
nodes_generator, message_func, reduce_func, apply_node_func
)
def prop_edges(
graph,
edges_generator,
message_func="default",
reduce_func="default",
apply_node_func="default",
):
"""Functional method for :func:`dgl.DGLGraph.prop_edges`.
Parameters
......@@ -54,14 +67,19 @@ def prop_edges(graph,
--------
dgl.DGLGraph.prop_edges
"""
graph.prop_edges(edges_generator, message_func, reduce_func, apply_node_func)
def prop_nodes_bfs(graph,
source,
message_func,
reduce_func,
reverse=False,
apply_node_func=None):
graph.prop_edges(
edges_generator, message_func, reduce_func, apply_node_func
)
def prop_nodes_bfs(
graph,
source,
message_func,
reduce_func,
reverse=False,
apply_node_func=None,
):
"""Message propagation using node frontiers generated by BFS.
Parameters
......@@ -83,10 +101,12 @@ def prop_nodes_bfs(graph,
--------
dgl.traversal.bfs_nodes_generator
"""
assert isinstance(graph, DGLHeteroGraph), \
'DGLGraph is deprecated, Please use DGLHeteroGraph'
assert len(graph.canonical_etypes) == 1, \
'prop_nodes_bfs only support homogeneous graph'
assert isinstance(
graph, DGLHeteroGraph
), "DGLGraph is deprecated, Please use DGLHeteroGraph"
assert (
len(graph.canonical_etypes) == 1
), "prop_nodes_bfs only support homogeneous graph"
# TODO(murphy): Graph traversal currently is only supported on
# CPP graphs. Move graph to CPU as a workaround,
# which should be fixed in the future.
......@@ -94,11 +114,10 @@ def prop_nodes_bfs(graph,
nodes_gen = [F.copy_to(frontier, graph.device) for frontier in nodes_gen]
prop_nodes(graph, nodes_gen, message_func, reduce_func, apply_node_func)
def prop_nodes_topo(graph,
message_func,
reduce_func,
reverse=False,
apply_node_func=None):
def prop_nodes_topo(
graph, message_func, reduce_func, reverse=False, apply_node_func=None
):
"""Message propagation using node frontiers generated by topological order.
Parameters
......@@ -118,10 +137,12 @@ def prop_nodes_topo(graph,
--------
dgl.traversal.topological_nodes_generator
"""
assert isinstance(graph, DGLHeteroGraph), \
'DGLGraph is deprecated, Please use DGLHeteroGraph'
assert len(graph.canonical_etypes) == 1, \
'prop_nodes_topo only support homogeneous graph'
assert isinstance(
graph, DGLHeteroGraph
), "DGLGraph is deprecated, Please use DGLHeteroGraph"
assert (
len(graph.canonical_etypes) == 1
), "prop_nodes_topo only support homogeneous graph"
# TODO(murphy): Graph traversal currently is only supported on
# CPP graphs. Move graph to CPU as a workaround,
# which should be fixed in the future.
......@@ -129,14 +150,17 @@ def prop_nodes_topo(graph,
nodes_gen = [F.copy_to(frontier, graph.device) for frontier in nodes_gen]
prop_nodes(graph, nodes_gen, message_func, reduce_func, apply_node_func)
def prop_edges_dfs(graph,
source,
message_func,
reduce_func,
reverse=False,
has_reverse_edge=False,
has_nontree_edge=False,
apply_node_func=None):
def prop_edges_dfs(
graph,
source,
message_func,
reduce_func,
reverse=False,
has_reverse_edge=False,
has_nontree_edge=False,
apply_node_func=None,
):
"""Message propagation using edge frontiers generated by labeled DFS.
Parameters
......@@ -162,15 +186,22 @@ def prop_edges_dfs(graph,
--------
dgl.traversal.dfs_labeled_edges_generator
"""
assert isinstance(graph, DGLHeteroGraph), \
'DGLGraph is deprecated, Please use DGLHeteroGraph'
assert len(graph.canonical_etypes) == 1, \
'prop_edges_dfs only support homogeneous graph'
assert isinstance(
graph, DGLHeteroGraph
), "DGLGraph is deprecated, Please use DGLHeteroGraph"
assert (
len(graph.canonical_etypes) == 1
), "prop_edges_dfs only support homogeneous graph"
# TODO(murphy): Graph traversal currently is only supported on
# CPP graphs. Move graph to CPU as a workaround,
# which should be fixed in the future.
edges_gen = trv.dfs_labeled_edges_generator(
graph.cpu(), source, reverse, has_reverse_edge, has_nontree_edge,
return_labels=False)
graph.cpu(),
source,
reverse,
has_reverse_edge,
has_nontree_edge,
return_labels=False,
)
edges_gen = [F.copy_to(frontier, graph.device) for frontier in edges_gen]
prop_edges(graph, edges_gen, message_func, reduce_func, apply_node_func)
python/dgl/random.py
View file @ 76bb5404
"""Python interfaces to DGL random number generators."""
import numpy as np
from ._ffi.function import _init_api
from . import backend as F
from . import ndarray as nd
from ._ffi.function import _init_api
__all__ = ["seed"]
__all__ = ['seed']
def seed(val):
"""Set the random seed of DGL.
......@@ -17,6 +18,7 @@ def seed(val):
"""
_CAPI_SetSeed(val)
def choice(a, size, replace=True, prob=None): # pylint: disable=invalid-name
"""An equivalent to :func:`numpy.random.choice`.
......@@ -57,7 +59,7 @@ def choice(a, size, replace=True, prob=None): # pylint: disable=invalid-name
samples : 1-D tensor
The generated random samples
"""
#TODO(minjie): support RNG as one of the arguments.
# TODO(minjie): support RNG as one of the arguments.
if isinstance(size, tuple):
num = np.prod(size)
else:
......@@ -74,7 +76,9 @@ def choice(a, size, replace=True, prob=None): # pylint: disable=invalid-name
prob = F.zerocopy_to_dgl_ndarray(prob)
bits = 64 # index array is in 64-bit
chosen_idx = _CAPI_Choice(int(num), int(population), prob, bool(replace), bits)
chosen_idx = _CAPI_Choice(
int(num), int(population), prob, bool(replace), bits
)
chosen_idx = F.zerocopy_from_dgl_ndarray(chosen_idx)
if F.is_tensor(a):
......@@ -87,4 +91,5 @@ def choice(a, size, replace=True, prob=None): # pylint: disable=invalid-name
else:
return chosen
_init_api('dgl.rng', __name__)
_init_api("dgl.rng", __name__)
python/dgl/readout.py
View file @ 76bb5404
"""Classes and functions for batching multiple graphs together."""
from __future__ import absolute_import
from .base import DGLError, dgl_warning
from . import backend as F
from .base import DGLError, dgl_warning
from .ops import segment
__all__ = ['readout_nodes', 'readout_edges',
'sum_nodes', 'sum_edges', 'mean_nodes', 'mean_edges',
'max_nodes', 'max_edges', 'softmax_nodes', 'softmax_edges',
'broadcast_nodes', 'broadcast_edges', 'topk_nodes', 'topk_edges']
def readout_nodes(graph, feat, weight=None, *, op='sum', ntype=None):
__all__ = [
"readout_nodes",
"readout_edges",
"sum_nodes",
"sum_edges",
"mean_nodes",
"mean_edges",
"max_nodes",
"max_edges",
"softmax_nodes",
"softmax_edges",
"broadcast_nodes",
"broadcast_edges",
"topk_nodes",
"topk_edges",
]
def readout_nodes(graph, feat, weight=None, *, op="sum", ntype=None):
"""Generate a graph-level representation by aggregating node features
:attr:`feat`.
......@@ -87,7 +100,8 @@ def readout_nodes(graph, feat, weight=None, *, op='sum', ntype=None):
x = x * graph.nodes[ntype].data[weight]
return segment.segment_reduce(graph.batch_num_nodes(ntype), x, reducer=op)
def readout_edges(graph, feat, weight=None, *, op='sum', etype=None):
def readout_edges(graph, feat, weight=None, *, op="sum", etype=None):
"""Sum the edge feature :attr:`feat` in :attr:`graph`, optionally
multiplies it by a edge :attr:`weight`.
......@@ -170,6 +184,7 @@ def readout_edges(graph, feat, weight=None, *, op='sum', etype=None):
x = x * graph.edges[etype].data[weight]
return segment.segment_reduce(graph.batch_num_edges(etype), x, reducer=op)
def sum_nodes(graph, feat, weight=None, *, ntype=None):
"""Syntax sugar for ``dgl.readout_nodes(graph, feat, weight, ntype=ntype, op='sum')``.
......@@ -177,7 +192,8 @@ def sum_nodes(graph, feat, weight=None, *, ntype=None):
--------
readout_nodes
"""
return readout_nodes(graph, feat, weight, ntype=ntype, op='sum')
return readout_nodes(graph, feat, weight, ntype=ntype, op="sum")
def sum_edges(graph, feat, weight=None, *, etype=None):
"""Syntax sugar for ``dgl.readout_edges(graph, feat, weight, etype=etype, op='sum')``.
......@@ -186,7 +202,8 @@ def sum_edges(graph, feat, weight=None, *, etype=None):
--------
readout_edges
"""
return readout_edges(graph, feat, weight, etype=etype, op='sum')
return readout_edges(graph, feat, weight, etype=etype, op="sum")
def mean_nodes(graph, feat, weight=None, *, ntype=None):
"""Syntax sugar for ``dgl.readout_nodes(graph, feat, weight, ntype=ntype, op='mean')``.
......@@ -195,7 +212,8 @@ def mean_nodes(graph, feat, weight=None, *, ntype=None):
--------
readout_nodes
"""
return readout_nodes(graph, feat, weight, ntype=ntype, op='mean')
return readout_nodes(graph, feat, weight, ntype=ntype, op="mean")
def mean_edges(graph, feat, weight=None, *, etype=None):
"""Syntax sugar for ``dgl.readout_edges(graph, feat, weight, etype=etype, op='mean')``.
......@@ -204,7 +222,8 @@ def mean_edges(graph, feat, weight=None, *, etype=None):
--------
readout_edges
"""
return readout_edges(graph, feat, weight, etype=etype, op='mean')
return readout_edges(graph, feat, weight, etype=etype, op="mean")
def max_nodes(graph, feat, weight=None, *, ntype=None):
"""Syntax sugar for ``dgl.readout_nodes(graph, feat, weight, ntype=ntype, op='max')``.
......@@ -213,7 +232,8 @@ def max_nodes(graph, feat, weight=None, *, ntype=None):
--------
readout_nodes
"""
return readout_nodes(graph, feat, weight, ntype=ntype, op='max')
return readout_nodes(graph, feat, weight, ntype=ntype, op="max")
def max_edges(graph, feat, weight=None, *, etype=None):
"""Syntax sugar for ``dgl.readout_edges(graph, feat, weight, etype=etype, op='max')``.
......@@ -222,7 +242,8 @@ def max_edges(graph, feat, weight=None, *, etype=None):
--------
readout_edges
"""
return readout_edges(graph, feat, weight, etype=etype, op='max')
return readout_edges(graph, feat, weight, etype=etype, op="max")
def softmax_nodes(graph, feat, *, ntype=None):
r"""Perform graph-wise softmax on the node features.
......@@ -283,6 +304,7 @@ def softmax_nodes(graph, feat, *, ntype=None):
x = graph.nodes[ntype].data[feat]
return segment.segment_softmax(graph.batch_num_nodes(ntype), x)
def softmax_edges(graph, feat, *, etype=None):
r"""Perform graph-wise softmax on the edge features.
......@@ -348,6 +370,7 @@ def softmax_edges(graph, feat, *, etype=None):
x = graph.edges[etype].data[feat]
return segment.segment_softmax(graph.batch_num_edges(etype), x)
def broadcast_nodes(graph, graph_feat, *, ntype=None):
"""Generate a node feature equal to the graph-level feature :attr:`graph_feat`.
......@@ -416,12 +439,15 @@ def broadcast_nodes(graph, graph_feat, *, ntype=None):
--------
broadcast_edges
"""
if (F.shape(graph_feat)[0] != graph.batch_size and graph.batch_size == 1):
dgl_warning('For a single graph, use a tensor of shape (1, *) for graph_feat.'
' The support of shape (*) will be deprecated.')
if F.shape(graph_feat)[0] != graph.batch_size and graph.batch_size == 1:
dgl_warning(
"For a single graph, use a tensor of shape (1, *) for graph_feat."
" The support of shape (*) will be deprecated."
)
graph_feat = F.unsqueeze(graph_feat, dim=0)
return F.repeat(graph_feat, graph.batch_num_nodes(ntype), dim=0)
def broadcast_edges(graph, graph_feat, *, etype=None):
"""Generate an edge feature equal to the graph-level feature :attr:`graph_feat`.
......@@ -487,17 +513,21 @@ def broadcast_edges(graph, graph_feat, *, etype=None):
--------
broadcast_nodes
"""
if (F.shape(graph_feat)[0] != graph.batch_size and graph.batch_size == 1):
dgl_warning('For a single graph, use a tensor of shape (1, *) for graph_feat.'
' The support of shape (*) will be deprecated.')
if F.shape(graph_feat)[0] != graph.batch_size and graph.batch_size == 1:
dgl_warning(
"For a single graph, use a tensor of shape (1, *) for graph_feat."
" The support of shape (*) will be deprecated."
)
graph_feat = F.unsqueeze(graph_feat, dim=0)
return F.repeat(graph_feat, graph.batch_num_edges(etype), dim=0)
READOUT_ON_ATTRS = {
'nodes': ('ndata', 'batch_num_nodes', 'number_of_nodes'),
'edges': ('edata', 'batch_num_edges', 'number_of_edges'),
"nodes": ("ndata", "batch_num_nodes", "number_of_nodes"),
"edges": ("edata", "batch_num_edges", "number_of_edges"),
}
def _topk_torch(keys, k, descending, x):
"""Internal function to take graph-wise top-k node/edge features according to
the rank given by keys, this function is PyTorch only.
......@@ -522,14 +552,18 @@ def _topk_torch(keys, k, descending, x):
A tensor with shape :math:`(batch, k)`.
"""
import torch as th
batch_size, max_len = x.shape[0], x.shape[1]
topk_indices = keys.topk(k, -1, largest=descending)[1] # (batch_size, k)
x = x.view((batch_size * max_len), -1)
shift = th.arange(0, batch_size, device=x.device).view(batch_size, 1) * max_len
shift = (
th.arange(0, batch_size, device=x.device).view(batch_size, 1) * max_len
)
topk_indices_ = topk_indices + shift
x = x[topk_indices_].view(batch_size, k, -1)
return th.masked_fill(x, th.isinf(x), 0), topk_indices
def _topk_on(graph, typestr, feat, k, descending, sortby, ntype_or_etype):
"""Internal function to take graph-wise top-k node/edge features of
field :attr:`feat` in :attr:`graph` ranked by keys at given
......@@ -577,24 +611,28 @@ def _topk_on(graph, typestr, feat, k, descending, sortby, ntype_or_etype):
_, batch_num_objs_attr, _ = READOUT_ON_ATTRS[typestr]
data = getattr(graph, typestr)[ntype_or_etype].data
if F.ndim(data[feat]) > 2:
raise DGLError('Only support {} feature `{}` with dimension less than or'
' equal to 2'.format(typestr, feat))
raise DGLError(
"Only support {} feature `{}` with dimension less than or"
" equal to 2".format(typestr, feat)
)
feat = data[feat]
hidden_size = F.shape(feat)[-1]
batch_num_objs = getattr(graph, batch_num_objs_attr)(ntype_or_etype)
batch_size = len(batch_num_objs)
length = max(max(F.asnumpy(batch_num_objs)), k)
fill_val = -float('inf') if descending else float('inf')
feat_ = F.pad_packed_tensor(feat, batch_num_objs, fill_val, l_min=k) # (batch_size, l, d)
fill_val = -float("inf") if descending else float("inf")
feat_ = F.pad_packed_tensor(
feat, batch_num_objs, fill_val, l_min=k
) # (batch_size, l, d)
if F.backend_name == 'pytorch' and sortby is not None:
if F.backend_name == "pytorch" and sortby is not None:
# PyTorch's implementation of top-K
keys = feat_[..., sortby] # (batch_size, l)
return _topk_torch(keys, k, descending, feat_)
else:
# Fallback to framework-agnostic implementation of top-K
if sortby is not None:
keys = F.squeeze(F.slice_axis(feat_, -1, sortby, sortby+1), -1)
keys = F.squeeze(F.slice_axis(feat_, -1, sortby, sortby + 1), -1)
order = F.argsort(keys, -1, descending=descending)
else:
order = F.argsort(feat_, 1, descending=descending)
......@@ -607,14 +645,18 @@ def _topk_on(graph, typestr, feat, k, descending, sortby, ntype_or_etype):
topk_indices_ = F.reshape(topk_indices, (-1,)) + shift
else:
feat_ = F.reshape(feat_, (-1,))
shift = F.repeat(F.arange(0, batch_size), k * hidden_size, -1) * length * hidden_size +\
F.cat([F.arange(0, hidden_size)] * batch_size * k, -1)
shift = F.repeat(
F.arange(0, batch_size), k * hidden_size, -1
) * length * hidden_size + F.cat(
[F.arange(0, hidden_size)] * batch_size * k, -1
)
shift = F.copy_to(shift, F.context(feat))
topk_indices_ = F.reshape(topk_indices, (-1,)) * hidden_size + shift
out = F.reshape(F.gather_row(feat_, topk_indices_), (batch_size, k, -1))
out = F.replace_inf_with_zero(out)
return out, topk_indices
def topk_nodes(graph, feat, k, *, descending=True, sortby=None, ntype=None):
"""Return a graph-level representation by a graph-wise top-k on
node features :attr:`feat` in :attr:`graph` by feature at index :attr:`sortby`.
......@@ -719,9 +761,16 @@ def topk_nodes(graph, feat, k, *, descending=True, sortby=None, ntype=None):
[3, 2, 0, 2, 2],
[2, 3, 2, 1, 3]]]))
"""
return _topk_on(graph, 'nodes', feat, k,
descending=descending, sortby=sortby,
ntype_or_etype=ntype)
return _topk_on(
graph,
"nodes",
feat,
k,
descending=descending,
sortby=sortby,
ntype_or_etype=ntype,
)
def topk_edges(graph, feat, k, *, descending=True, sortby=None, etype=None):
"""Return a graph-level representation by a graph-wise top-k
......@@ -827,6 +876,12 @@ def topk_edges(graph, feat, k, *, descending=True, sortby=None, etype=None):
[3, 2, 0, 2, 2],
[2, 3, 2, 1, 3]]]))
"""
return _topk_on(graph, 'edges', feat, k,
descending=descending, sortby=sortby,
ntype_or_etype=etype)
return _topk_on(
graph,
"edges",
feat,
k,
descending=descending,
sortby=sortby,
ntype_or_etype=etype,
)
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment