CuGraphSAGEConv

class dgl.nn.pytorch.conv.CuGraphSAGEConv(in_feats, out_feats, aggregator_type='mean', feat_drop=0.0, bias=True)[source]

基类: CuGraphBaseConv

一个加速的 GraphSAGE 层，来自论文 Inductive Representation Learning on Large Graphs，它利用了 cugraph-ops 中高度优化的聚合原语

\[ \begin{align}\begin{aligned}h_{\mathcal{N}(i)}^{(l+1)} &= \mathrm{aggregate} \left(\{h_{j}^{l}, \forall j \in \mathcal{N}(i) \}\right)\\h_{i}^{(l+1)} &= W \cdot \mathrm{concat} (h_{i}^{l}, h_{\mathcal{N}(i)}^{(l+1)})\end{aligned}\end{align} \]

此模块依赖于 pylibcugraphops 包，可以通过 conda install -c nvidia pylibcugraphops=23.04 进行安装。pylibcugraphops 23.04 需要 python 3.8.x 或 3.10.x 版本。

注意

这是一项实验性功能。

参数:

in_feats (int) – 输入特征大小。
out_feats (int) – 输出特征大小。
aggregator_type (str) – 使用的聚合器类型（mean, sum, min, max）。
feat_drop (float) – 特征上的 Dropout 率，默认为: 0。
bias (bool) – 如果为 True，则在输出中添加一个可学习的偏置项。默认为: True。

示例

>>> import dgl
>>> import torch
>>> from dgl.nn import CuGraphSAGEConv
>>> device = 'cuda'
>>> g = dgl.graph(([0,1,2,3,2,5], [1,2,3,4,0,3])).to(device)
>>> g = dgl.add_self_loop(g)
>>> feat = torch.ones(6, 10).to(device)
>>> conv = CuGraphSAGEConv(10, 2, 'mean').to(device)
>>> res = conv(g, feat)
>>> res
tensor([[-1.1690,  0.1952],
        [-1.1690,  0.1952],
        [-1.1690,  0.1952],
        [-1.1690,  0.1952],
        [-1.1690,  0.1952],
        [-1.1690,  0.1952]], device='cuda:0', grad_fn=<AddmmBackward0>)

forward(g, feat, max_in_degree=None)[source]

前向计算。

参数:

g (DGLGraph) – 图。
feat (torch.Tensor) – 节点特征。形状: \((N, D_{in})\)。
max_in_degree (int) – 目标节点的最大入度。仅当 g 是一个 DGLBlock（即二分图）时有效。当 g 由邻居采样器生成时，该值应设置为相应的 fanout。如果未给定，将即时计算 max_in_degree。

返回:

输出节点特征。形状: \((N, D_{out})\)。

返回类型:

torch.Tensor

reset_parameters()[source]: 重新初始化可学习参数。