DGL-dgl.function
这个子包承载了 DGL 提供的所有内置函数。内置函数是 DGL 推荐的表达:消息传递计算(即通过update_all())或从节点特征计算边缘特征(即通过apply_edges())。内置函数以符号方式描述节点方式和边方式计算,无需任何实际计算,因此 DGL 可以分析并将它们映射到高效的低级内核。这里有些例子:
import dgl
import dgl.function as fn
import torch as th
g = ... # create a DGLGraph
g.ndata['h'] = th.randn((g.num_nodes(), 10)) # each node has feature size 10
g.edata['w'] = th.randn((g.num_edges(), 1)) # each edge has feature size 1
# collect features from source nodes and aggregate them in destination nodes
g.update_all(fn.copy_u('h', 'm'), fn.sum('m', 'h_sum'))
# multiply source node features with edge weights and aggregate them in destination nodes
g.update_all(fn.u_mul_e('h', 'w', 'm'), fn.max('m', 'h_max'))
# compute edge embedding by multiplying source and destination node embeddings
g.apply_edges(fn.u_mul_v('h', 'h', 'w_new'))
fn.copy_u
, fn.u_mul_e
,fn.u_mul_v
是内置的消息函数,而fn.sum
和fn.max
是内置的 reduce 函数。DGL 的约定是使用u,v 和e分别表示源节点、目的节点和边。例如,copy_u
告诉 DGL 将源节点数据复制为消息; u_mul_e
告诉 DGL 将源节点特征与边特征相乘。
要定义一元消息函数(例如copy_u
),请指定一个输入特征名称和一个输出消息名称。要定义二进制消息函数(例如u_mul_e
),请指定两个输入特征名称和一个输出消息名称。在计算过程中,消息函数将读取给定名称下的数据,执行计算,并使用输出名称返回输出。例如,上面与下面的用户定义函数相同:fn.u_mul_e('h', 'w', 'm')
def udf_u_mul_e(edges):
return {'m' : edges.src['h'] * edges.data['w']}
要定义一个 reduce
函数,需要指定一个输入消息名称和一个输出节点特征名称。例如,上面与下面的用户定义函数相同:fn.max('m', 'h_max')
def udf_max(nodes):
return {'h_max' : th.max(nodes.mailbox['m'], 1)[0]}
所有二进制消息函数都支持广播,这是一种将逐元素操作扩展到具有不同形状的张量输入的机制。DGL 通常遵循NumPy 和PyTorch的标准广播语义。下面是一些例子:
import dgl
import dgl.function as fn
import torch as th
g = ... # create a DGLGraph
# case 1
g.ndata['h'] = th.randn((g.num_nodes(), 10))
g.edata['w'] = th.randn((g.num_edges(), 1))
# OK, valid broadcasting between feature shapes (10,) and (1,)
g.update_all(fn.u_mul_e('h', 'w', 'm'), fn.sum('m', 'h_new'))
g.ndata['h_new'] # shape: (g.num_nodes(), 10)
# case 2
g.ndata['h'] = th.randn((g.num_nodes(), 5, 10))
g.edata['w'] = th.randn((g.num_edges(), 10))
# OK, valid broadcasting between feature shapes (5, 10) and (10,)
g.update_all(fn.u_mul_e('h', 'w', 'm'), fn.sum('m', 'h_new'))
g.ndata['h_new'] # shape: (g.num_nodes(), 5, 10)
# case 3
g.ndata['h'] = th.randn((g.num_nodes(), 5, 10))
g.edata['w'] = th.randn((g.num_edges(), 5))
# NOT OK, invalid broadcasting between feature shapes (5, 10) and (5,)
# shapes are aligned from right
g.update_all(fn.u_mul_e('h', 'w', 'm'), fn.sum('m', 'h_new'))
# case 3
g.ndata['h1'] = th.randn((g.num_nodes(), 1, 10))
g.ndata['h2'] = th.randn((g.num_nodes(), 5, 1))
# OK, valid broadcasting between feature shapes (1, 10) and (5, 1)
g.apply_edges(fn.u_add_v('h1', 'h2', 'x')) # apply_edges also supports broadcasting
g.edata['x'] # shape: (g.num_edges(), 5, 10)
# case 4
g.ndata['h1'] = th.randn((g.num_nodes(), 1, 10, 128))
g.ndata['h2'] = th.randn((g.num_nodes(), 5, 1, 128))
# OK, u_dot_v supports broadcasting but requires the last dimension to match
g.apply_edges(fn.u_dot_v('h1', 'h2', 'x'))
g.edata['x'] # shape: (g.num_edges(), 5, 10, 1)
消息函数
-
copy_src
(来源,输出)
使用源节点功能计算消息的内置消息函数。 -
copy_edge
(边缘,出)
使用边缘特征计算消息的内置消息函数。 -
src_mul_edge
(源、边、出)
内置消息函数,通过在 src 特征和边缘特征之间执行二元运算 mul 来计算消息。 -
copy_u
(你,出去)
使用源节点功能计算消息的内置消息函数。 -
copy_e
(e,出来)
使用边缘特征计算消息的内置消息函数。 -
u_add_v(lhs_field, rhs_field, out)
内置消息函数,如果特征具有相同的形状,则通过在 u 和 v 的特征之间执行逐元素相加来计算边上的消息;否则,它首先将特征广播到一个新的形状并执行逐元素操作。 -
u_sub_v(lhs_field, rhs_field, out)
内置消息函数,如果特征具有相同的形状,则通过在 u 和 v 的特征之间执行逐元素子运算来计算边上的消息;否则,它首先将特征广播到一个新的形状并执行逐元素操作。 -
u_mul_v(lhs_field, rhs_field, out)
内置消息函数,如果特征具有相同的形状,则通过在 u 和 v 的特征之间执行逐元素乘法来计算边上的消息;否则,它首先将特征广播到一个新的形状并执行逐元素操作。 -
u_div_v(lhs_field, rhs_field, out)
内置消息函数,如果特征具有相同的形状,则通过在 u 和 v 的特征之间执行逐元素 div 来计算边上的消息;否则,它首先将特征广播到一个新的形状并执行逐元素操作。 -
u_add_e(lhs_field, rhs_field, out)
内置消息函数,如果特征具有相同的形状,则通过在 u 和 e 的特征之间执行逐元素相加来计算边上的消息;否则,它首先将特征广播到一个新的形状并执行逐元素操作。 -
u_sub_e(lhs_field, rhs_field, out)
内置消息函数,如果特征具有相同的形状,则通过在 u 和 e 的特征之间执行逐元素子运算来计算边上的消息;否则,它首先将特征广播到一个新的形状并执行逐元素操作。 -
u_mul_e(lhs_field, rhs_field, out)
内置消息函数,如果特征具有相同的形状,则通过在 u 和 e 的特征之间执行逐元素乘法来计算边上的消息;否则,它首先将特征广播到一个新的形状并执行逐元素操作。 -
u_div_e(lhs_field, rhs_field, out)
内置消息函数,如果特征具有相同的形状,则通过在 u 和 e 的特征之间执行逐元素 div 来计算边上的消息;否则,它首先将特征广播到一个新的形状并执行逐元素操作。 -
v_add_u(lhs_field, rhs_field, out)
内置消息函数,如果特征具有相同的形状,则通过在 v 和 u 的特征之间执行逐元素相加来计算边上的消息;否则,它首先将特征广播到一个新的形状并执行逐元素操作。 -
v_sub_u(lhs_field, rhs_field, out)
内置消息函数,如果特征具有相同的形状,则通过在 v 和 u 的特征之间执行逐元素子运算来计算边上的消息;否则,它首先将特征广播到一个新的形状并执行逐元素操作。 -
v_mul_u(lhs_field, rhs_field, out)
内置消息函数,如果特征具有相同的形状,则通过在 v 和 u 的特征之间执行逐元素乘法来计算边上的消息;否则,它首先将特征广播到一个新的形状并执行逐元素操作。 -
v_div_u(lhs_field, rhs_field, out)
内置消息函数,如果特征具有相同的形状,则通过在 v 和 u 的特征之间执行逐元素 div 来计算边上的消息;否则,它首先将特征广播到一个新的形状并执行逐元素操作。 -
v_add_e(lhs_field, rhs_field, out)
内置消息函数,如果特征具有相同的形状,则通过在 v 和 e 的特征之间执行逐元素相加来计算边上的消息;否则,它首先将特征广播到一个新的形状并执行逐元素操作。 -
v_sub_e(lhs_field, rhs_field, out)
内置消息函数,如果特征具有相同的形状,则通过在 v 和 e 的特征之间执行逐元素子运算来计算边上的消息;否则,它首先将特征广播到一个新的形状并执行逐元素操作。 -
v_mul_e(lhs_field, rhs_field, out)
内置消息函数,如果特征具有相同的形状,则通过在 v 和 e 的特征之间执行逐元素乘法来计算边上的消息;否则,它首先将特征广播到一个新的形状并执行逐元素操作。 -
v_div_e(lhs_field, rhs_field, out)
内置消息函数,如果特征具有相同的形状,则通过在 v 和 e 的特征之间执行逐元素 div 来计算边上的消息;否则,它首先将特征广播到一个新的形状并执行逐元素操作。 -
e_add_u(lhs_field, rhs_field, out)
内置消息函数,如果特征具有相同的形状,则通过在 e 和 u 的特征之间执行逐元素相加来计算边上的消息;否则,它首先将特征广播到一个新的形状并执行逐元素操作。
e_sub_u(lhs_field, rhs_field, out)
内置消息函数,如果特征具有相同的形状,则通过在 e 和 u 的特征之间执行逐元素子运算来计算边上的消息;否则,它首先将特征广播到一个新的形状并执行逐元素操作。
-
e_mul_u(lhs_field, rhs_field, out)
内置消息函数,如果特征具有相同的形状,则通过在 e 和 u 的特征之间执行逐元素乘法来计算边上的消息;否则,它首先将特征广播到一个新的形状并执行逐元素操作。 -
e_div_u(lhs_field, rhs_field, out)
内置消息函数,如果特征具有相同的形状,则通过在 e 和 u 的特征之间执行逐元素 div 来计算边上的消息;否则,它首先将特征广播到一个新的形状并执行逐元素操作。 -
e_add_v(lhs_field, rhs_field, out)
内置消息函数,如果特征具有相同的形状,则通过在 e 和 v 的特征之间执行逐元素相加来计算边上的消息;否则,它首先将特征广播到一个新的形状并执行逐元素操作。 -
e_sub_v(lhs_field, rhs_field, out)
内置消息函数,如果特征具有相同的形状,则通过在 e 和 v 的特征之间执行逐元素子运算来计算边上的消息;否则,它首先将特征广播到一个新的形状并执行逐元素操作。 -
e_mul_v(lhs_field, rhs_field, out)
内置消息函数,如果特征具有相同的形状,则通过在 e 和 v 的特征之间执行逐元素乘法来计算边上的消息;否则,它首先将特征广播到一个新的形状并执行逐元素操作。 -
e_div_v(lhs_field, rhs_field, out)
内置消息函数,如果特征具有相同的形状,则通过在 e 和 v 的特征之间执行逐元素 div 来计算边上的消息;否则,它首先将特征广播到一个新的形状并执行逐元素操作。 -
u_dot_v(lhs_field, rhs_field, out)
内置消息函数,如果特征具有相同的形状,则通过在 u 和 v 的特征之间执行逐元素点来计算边上的消息;否则,它首先将特征广播到一个新的形状并执行逐元素操作。 -
u_dot_e(lhs_field, rhs_field, out)
内置消息函数,如果特征具有相同的形状,则通过在 u 和 e 的特征之间执行逐元素点来计算边上的消息;否则,它首先将特征广播到一个新的形状并执行逐元素操作。
v_dot_e(lhs_field, rhs_field, out)
内置消息函数,如果特征具有相同的形状,则通过在 v 和 e 的特征之间执行逐元素点来计算边上的消息;否则,它首先将特征广播到一个新的形状并执行逐元素操作。
-
v_dot_u(lhs_field, rhs_field, out)
内置消息函数,如果特征具有相同的形状,则通过在 v 和 u 的特征之间执行逐元素点来计算边上的消息;否则,它首先将特征广播到一个新的形状并执行逐元素操作。 -
e_dot_u(lhs_field, rhs_field, out)
内置消息函数,如果特征具有相同的形状,则通过在 e 和 u 的特征之间执行逐元素点来计算边上的消息;否则,它首先将特征广播到一个新的形状并执行逐元素操作。 -
e_dot_v(lhs_field, rhs_field, out)
内置消息函数,如果特征具有相同的形状,则通过在 e 和 v 的特征之间执行逐元素点来计算边上的消息;否则,它首先将特征广播到一个新的形状并执行逐元素操作。
Reduce函数
-
sum
(消息,输出)
内置 reduce 函数,按总和聚合消息。 -
max
(消息,输出)
内置的 reduce 函数可以最大程度地聚合消息。 -
min
(消息,输出)
内置 reduce 函数,按最小值聚合消息。 -
mean
(消息,输出)
内置 reduce 函数,按均值聚合消息。