DGL-dgl.function

发呆的比目鱼

已于 2023-01-12 11:14:15 修改

阅读量986

点赞数

分类专栏：图神经网络文章标签：人工智能

于 2023-01-12 09:41:48 首次发布

本文链接：https://blog.csdn.net/weixin_42486623/article/details/128654034

版权

图神经网络专栏收录该内容

62 篇文章 19 订阅

订阅专栏

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 的特征之间执行逐元素点来计算边上的消息；否则，它首先将特征广播到一个新的形状并执行逐元素操作。