0

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

我想GatedGraphConv在可能具有许多不同边类型的图上使用一个模块，并且想知道对于apply_edges()下面的任何给定调用，该调用是否只是在 GPU 上调度任务，或者该调用是否阻塞了循环的剩余迭代，直到计算完成后？

这是forward()方法的主循环GatedGraphConv（可在https://docs.dgl.ai/_modules/dgl/nn/pytorch/conv/gatedgraphconv.html#GatedGraphConv找到）：

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)

我已经研究过可能在单个 GPU 上使用一组 cuda 流，其中每个流都会被赋予自己的调用apply_edges()（每种不同的边缘类型使用不同的字符串"W_e*h"，因此它们不会相互冲突），然后在之前同步它们的调用update_all()，但我想到这可能是矫枉过正，也许调用apply_edges()只是在 GPU 上安排一个任务，然后 GPU 可以自由地以它认为合适的任何方式完成所有预定的任务（包括运行如果可以的话，同时执行多个任务）。

谢谢。

4

0 回答 0