In the doc of forward func for SAGEConv:
Returns: The output feature of shape (π,π·ππ’π‘) where π·ππ’π‘ is size of output feature.
N is the number of nodes in the graph.
However, this may be wrong for block sampled by dataloader:
import dgl
import torch as th
from dgl.nn import SAGEConv
g = dgl.graph(([0, 2, 0, 1, 3, 4], [1, 1, 2, 0, 1, 0]))
train_nid = th.tensor([1, 0, 2, 3], dtype=th.int64)
sampler = dgl.dataloading.MultiLayerNeighborSampler([3])
dataloader = dgl.dataloading.NodeDataLoader(
g, train_nid, sampler,
batch_size=2, shuffle=True, drop_last=False, num_workers=1)
conv = SAGEConv(10, 2, 'pool')
feat = th.nn.Embedding(5, 10)
for input_nodes, output_nodes, blocks in dataloader:
print(blocks[0])
print(conv(blocks[0], feat(input_nodes)))
The returned shape should be (π_dst,π·ππ’π‘).