DGL官方教程--API--dgl.function

参考：https://docs.dgl.ai/en/latest/api/python/function.html

dgl.function 

在DGL中，消息传递由两个API表示：

• send(edges,message_func)用于计算沿给定边缘的消息。
• recv(nodes,reduce_func) 用于收集传入消息，执行聚合等。

尽管两阶段抽象可以覆盖消息传递范式中定义的所有模型，但由于需要存储显式消息，所以效率低下。有关更多详细信息和性能结果，请参见DGL 博客。

我们的解决方案（也在博客文章中解释了）是将两个阶段融合到一个内核中，这样就不会生成和存储任何显式消息。为此，我们建议使用built-in message and reduce functions，以便DGL可以分析并将它们映射到融合的专用内核。以下是一些示例（使用PyTorch语法）。

import dgl
import dgl.function as fn
import torch as th
g = ... # 创建一个DGL图
g.ndata['h'] = th.randn((g.number_of_nodes(), 10)) # 每个节点的特征size为10。
g.edata['w'] = th.randn((g.number_of_edges(), 1))  # 每条边的特征size为1
# 从源节点收集特性，并将它们聚合到目标节点中。
g.update_all(fn.copy_u('h', 'm'), fn.sum('m', 'h_sum'))
# 具有边缘权重的多源节点特征，并将它们聚合到目标节点中。
g.update_all(fn.u_mul_e('h', 'w', 'm'), fn.max('m', 'h_max'))
# 通过多源和目标节点嵌入计算边缘嵌入
g.apply_edges(fn.u_mul_v('h', 'h', 'w_new'))

fn.copy_u，fn.u_mul_e，fn.u_mul_v是内置的message functions，而fn.sum 且fn.max是内置的 reduce functions。我们使用u，v和e代表源节点，目的节点和边。因此，copy_u将源节点数据复制为消息，u_mul_e是将源节点特征与边缘特征相乘。

要定义一元消息功能（例如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 function，需要指定一个输入消息名称和一个输出节点特征名称。例如，以上fn.max('m','h_max')与以下用户定义函数相同：

def udf_max(nodes):
   return {'h_max' : th.max(nodes.mailbox['m'], 1)[0]}

二元消息功能支持广播，这意味着张量参数可以自动扩展为相等大小。支持的广播语法是标准的，并且与NumPy 和PyTorch相同。如果您不熟悉广播，请参阅链接的主题以了解更多信息。在上面的示例中，fn.u_mul_e由于节点特征'h'和边缘特征'w'的形状不同，将自动执行广播乘法，但是可以广播它们。

DGL的所有内置功能都支持CPU和GPU以及向后计算，因此它们可以在任何autograd系统中使用。另外，内置函数不仅可以在使用update_all 或apply_edges如实施例中示出，但无论是必需的消息，并减少功能（例如pull，push，send_and_recv）。

这是所有DGL内置函数的备忘单。

类别	职能	备忘录
一元message function	copy_u
	copy_e
	copy_src	copy_u的别名
	copy_edge	copy_e的别名
二元message function	u_add_v，u_sub_v，u_mul_v，u_div_v，u_dot_v
	u_add_e，u_sub_e，u_mul_e，u_div_e，u_dot_e
	v_add_u，v_sub_u，v_mul_u，v_div_u，v_dot_u
	v_add_e，v_sub_e，v_mul_e，v_div_e，v_dot_e
	e_add_u，e_sub_u，e_mul_u，e_div_u，e_dot_u
	e_add_v，e_sub_v，e_mul_v，e_div_v，e_dot_v
	src_mul_edge	u_mul_e的别名
Reduce function	max
	min
	sum
	prod
	mean

消息功能

copy_src(src, out)	内置消息功能，使用源节点特征计算消息。
copy_edge(edge, out)	内置消息功能，可使用边缘特征计算消息。
src_mul_edge(src, edge, out)	内置消息功能，通过执行src特征和边缘特性之间的二元操作来计算消息。
copy_u(u, out)	内置消息功能，可使用源节点特征计算消息。
copy_e(e, out)	内置消息功能，可使用边特征计算消息。
u_add_v(lhs_field, rhs_field, out)	内置消息功能，如果特征形状相同，则通过在src和dst特征之间执行逐元素加法来计算边缘上的消息；否则，它将首先将特征广播到新形状并执行逐元素操作。
u_sub_v(lhs_field, rhs_field, out)	内置消息功能，如果特征形状相同，则通过在src和dst特征之间执行逐元素减法来计算边缘上的消息；否则，它将首先将特征广播到新形状并执行逐元素操作。
u_mul_v(lhs_field, rhs_field, out)	内置消息功能，如果特征形状相同，则通过在src和dst特征之间执行逐元素乘法来计算边缘上的消息；否则，它将首先将特征广播到新形状并执行逐元素操作。
u_div_v(lhs_field, rhs_field, out	内置消息功能，如果特征形状相同，则通过在src和dst特征之间执行逐元素除法来计算边缘上的消息；否则，它将首先将特征广播到新形状并执行逐元素操作。
u_add_e(lhs_field, rhs_field, out)	内置消息功能，如果特征形状相同，则通过在src和edge特征之间执行逐元素加法来计算边缘上的消息；否则，它将首先将特征广播到新形状并执行逐元素操作。
u_sub_e(lhs_field, rhs_field, out)                ）	内置消息功能，如果特征形状相同，则通过在src和edge特征之间执行逐元素减法来计算边缘上的消息；否则，它将首先将特征广播到新形状并执行逐元素操作。
u_mul_e(lhs_field, rhs_field, out)	内置消息功能，如果特征形状相同，则通过在src和edge特征之间执行逐元素乘法来计算边缘上的消息；否则，它将首先将特征广播到新形状并执行逐元素操作。
u_div_e(lhs_field, rhs_field, out)	内置消息功能，如果特征形状相同，则通过在src和edge特征之间执行逐元素除法来计算边缘上的消息；否则，它将首先将特征广播到新形状并执行逐元素操作。
v_add_u(lhs_field, rhs_field, out)	内置消息功能，如果特征形状相同，则通过在dst和src特征之间执行逐元素加法来计算边缘上的消息；否则，它将首先将特征广播到新形状并执行逐元素操作。
v_sub_u(lhs_field, rhs_field, out)	内置消息功能，如果特征形状相同，则通过在dst和src特征之间执行逐元素减法来在边缘上计算消息；否则，它将首先将特征广播到新形状并执行逐元素操作。
v_mul_u(lhs_field, rhs_field, out)	内置消息功能，如果特征形状相同，则通过在dst和src特征之间执行逐元素乘法来计算边缘上的消息；否则，它将首先将特征广播到新形状并执行逐元素操作。
v_div_u(lhs_field, rhs_field, out)	内置消息功能，如果特征形状相同，则通过在dst和src特征之间执行逐元素除法来计算边缘上的消息；否则，它将首先将特征广播到新形状并执行逐元素操作。
v_add_e(lhs_field, rhs_field, out)	内置消息功能，如果特征形状相同，则通过在dst和edge特征之间执行逐元素加法来计算边缘上的消息；否则，它将首先将特征广播到新形状并执行逐元素操作。
v_sub_e(lhs_field, rhs_field, out)	内置消息功能，如果特征形状相同，则通过在dst和edge特征之间执行逐元素减法来计算边缘上的消息；否则，它将首先将特征广播到新形状并执行逐元素操作。
v_mul_e(lhs_field, rhs_field, out)	内置消息功能，如果特征形状相同，则通过在dst和edge特征之间执行逐元素乘法来计算边缘上的消息；否则，它将首先将特征广播到新形状并执行逐元素操作。
v_div_e(lhs_field, rhs_field, out)	内置消息功能，如果特征形状相同，则通过在dst和edge特征之间执行逐元素除法来计算边缘上的消息；否则，它将首先将特征广播到新形状并执行逐元素操作。
e_add_u(lhs_field, rhs_field, out)	内置消息功能，如果特征形状相同，则通过在edge和src的特征之间执行逐元素加法来计算边缘上的消息；否则，它将首先将特征广播到新形状并执行逐元素操作。
e_sub_u(lhs_field, rhs_field, out)	内置消息功能，如果特征形状相同，则通过在edge和src的特征之间执行逐元素减法来计算边缘上的消息；否则，它将首先将特征广播到新形状并执行逐元素操作。
e_mul_u(lhs_field, rhs_field, out)	内置消息功能，如果特征形状相同，则通过在edge和src的特征之间执行逐元素乘法来计算边缘上的消息；否则，它将首先将特征广播到新形状并执行逐元素操作。
e_div_u(lhs_field, rhs_field, out)	内置消息功能，如果特征形状相同，则通过在edge和src的特征之间执行逐元素除法来计算边缘上的消息；否则，它将首先将特征广播到新形状并执行逐元素操作。
e_add_v(lhs_field, rhs_field, out)	内置消息功能，如果特征形状相同，则通过在edge和dst的特征之间执行逐元素加法来计算边缘上的消息；否则，它将首先将特征广播到新形状并执行逐元素操作。
e_sub_v(lhs_field, rhs_field, out)	内置消息功能，如果特征形状相同，则通过在edge和dst的特征之间执行逐元素减法来计算边缘上的消息；否则，它将首先将特征广播到新形状并执行逐元素操作。
e_mul_v(lhs_field, rhs_field, out)	内置消息功能，如果特征形状相同，则通过在edge和dst的特征之间执行逐元素乘法来计算边缘上的消息；否则，它将首先将特征广播到新形状并执行逐元素操作。
e_div_v(lhs_field, rhs_field, out)	内置消息功能，如果特征形状相同，则通过在edge和dst的特征之间执行逐元素除法来计算边缘上的消息；否则，它将首先将特征广播到新形状并执行逐元素操作。
u_dot_v(lhs_field, rhs_field, out)	内置消息功能，如果特征形状相同，则通过在src和dst特征之间执行逐元素点乘来计算边缘上的消息；否则，它将首先将特征广播到新形状并执行逐元素操作。
u_dot_e(lhs_field, rhs_field, out)	内置消息功能，如果特征形状相同，则通过在src和edge的特征之间执行逐元素点乘来计算边缘上的消息；否则，它将首先将特征广播到新形状并执行逐元素操作。
v_dot_e(lhs_field, rhs_field, out)	内置消息功能，如果特征形状相同，则通过在dst和edge的特征之间执行逐元素点乘计算边缘上的消息；否则，它将首先将特征广播到新形状并执行逐元素操作。
v_dot_u(lhs_field, rhs_field, out)	内置消息功能，如果特征形状相同，则通过在dst和src特征之间执行逐元素点乘计算边缘上的消息；否则，它将首先将特征广播到新形状并执行逐元素操作。
e_dot_u(lhs_field, rhs_field, out)	内置消息功能，如果特征形状相同，通过在edge和src的特征之间执行逐元素点乘计算边缘上的消息；否则，它将首先将特征广播到新形状并执行逐元素操作。
e_dot_v(lhs_field, rhs_field, out)	内置消息功能，如果特征形状相同，则通过在edge和dst的特征之间执行逐元素点乘计算边缘上的消息；否则，它将首先将特征广播到新形状并执行逐元素操作。

约简功能

sum(msg, out)	内置的reduce函数，按sum来汇总消息。
max(msg, out)	内置的reduce函数，按max来汇总消息。
min(msg, out)	内置的reduce函数，按min来汇总消息。
prod(msg, out)	内置的reduce函数，按prod来汇总消息。
mean(msg, out)	内置的reduce函数，按mean来汇总消息。

 

dgl.function.sum 

dgl.function.sum(msg,out)

内置的reduce函数，按sum来汇总消息。

参数：	msg（str）–消息字段。 out（str）–输出节点特征字段。

例子

>>> import dgl
>>> reduce_func = dgl.function.sum('m', 'h')

上面的示例等效于以下用户定义的函数（如果使用PyTorch）：

>>> import torch
>>> def reduce_func(nodes):
>>>     return {'h': torch.sum(nodes.mailbox['m'], dim=1)}

dgl.function.max 

dgl.function.max(msg,out)

内置的reduce函数，按max来汇总消息。

参数：	msg（str）–消息字段。 out（str）–输出节点特征字段。

例子

>>> import dgl
>>> reduce_func = dgl.function.max('m', 'h')

 

上面的示例等效于以下用户定义的函数（如果使用PyTorch）：

>>> import torch
>>> def reduce_func(nodes):
>>>     return {'h': torch.max(nodes.mailbox['m'], dim=1)}

dgl.function.min 

dgl.function.min(msg,out)

内置的reduce函数，按min来汇总消息。

参数：	msg（str）–消息字段。 out（str）–输出节点特征字段。

例子

>>> import dgl
>>> reduce_func = dgl.function.min('m', 'h')

 

上面的示例等效于以下用户定义的函数（如果使用PyTorch）：

>>> import torch
>>> def reduce_func(nodes):
>>>     return {'h': torch.min(nodes.mailbox['m'], dim=1)}

 

dgl.function.prod 

dgl.function.prod(msg,out)

内置的reduce函数，按prod来汇总消息。

参数：	msg（str）–消息字段。 out（str）–输出节点特征字段。

例子

>>> import dgl
>>> reduce_func = dgl.function.prod('m', 'h')

 

上面的示例等效于以下用户定义的函数（如果使用PyTorch）：

>>> import torch
>>> def reduce_func(nodes):
>>>     return {'h': torch.prod(nodes.mailbox['m'], dim=1)}

 

dgl.function.mean 

dgl.function.mean(msg,out)

内置的reduce函数，按mean来汇总消息。

参数：	msg（str）–消息字段。 out（str）–输出节点特征字段。

例子

>>> import dgl
>>> reduce_func = dgl.function.mean('m', 'h')

上面的示例等效于以下用户定义的函数（如果使用PyTorch）：

>>> import torch
>>> def reduce_func(nodes):
>>>     return {'h': torch.mean(nodes.mailbox['m'], dim=1)}