Pytorch 关于串行或并行执行DGL消息传递和聚合命令的问题

我的问题是关于在GPU上执行来自DGL的消息传递和聚合命令的顺序

我想在一个可能有许多不同边类型的图上使用一个

GatedGraphConv

模块，我想知道，对于下面的

apply_edges（）

的任何给定调用，调用是否只是在GPU上调度任务，或者调用是否在完成计算之前阻止循环的剩余迭代

以下是

GatedGraphConv

的

forward（）

方法的主循环（位于）：

我已经研究过可能在单个GPU上使用一组cuda流，其中每个cuda流都将被赋予自己对

apply_edges（）

（对每个不同的边缘类型使用不同的字符串

“W_e*h”

，这样它们就不会相互冲突），然后在调用

update_all（）

之前对它们进行同步，但我突然想到，这可能有些过分，调用

apply_edges（）

可能只是为了在GPU上安排一个任务，然后GPU可以自由地以它认为合适的任何方式完成所有计划的任务（包括同时运行多个任务，如果可以的话）

谢谢

for _ in range(self._n_steps):
    graph.ndata['h'] = feat
    for i in range(self._n_etypes):
        eids = th.nonzero(etypes == i, as_tuple=False).view(-1).type(graph.idtype)
        if len(eids) > 0:
            graph.apply_edges(
                lambda edges: {'W_e*h': self.linears[i](edges.src['h'])},
                eids
            )
    graph.update_all(fn.copy_e('W_e*h', 'm'), fn.sum('m', 'a'))
    a = graph.ndata.pop('a') # (N, D)
    feat = self.gru(a, feat)