wehub-resource-sync
118 lines
4.4 KiB
Python
118 lines
4.4 KiB
Python
"""Graphbolt cooperative convolution."""
|
|
from typing import Dict, Union
|
|
|
|
import torch
|
|
|
|
from ..sampled_subgraph import SampledSubgraph
|
|
from ..subgraph_sampler import all_to_all, convert_to_hetero, revert_to_homo
|
|
|
|
__all__ = ["CooperativeConvFunction", "CooperativeConv"]
|
|
|
|
|
|
class CooperativeConvFunction(torch.autograd.Function):
|
|
"""Cooperative convolution operation from Cooperative Minibatching.
|
|
|
|
Implements the `all-to-all` message passing algorithm
|
|
in Cooperative Minibatching, which was initially proposed in
|
|
`Deep Graph Library PR#4337<https://github.com/dmlc/dgl/pull/4337>`__ and
|
|
was later first fully described in
|
|
`Cooperative Minibatching in Graph Neural Networks
|
|
<https://arxiv.org/abs/2310.12403>`__.
|
|
Cooperation between the GPUs eliminates duplicate work performed across the
|
|
GPUs due to the overlapping sampled k-hop neighborhoods of seed nodes when
|
|
performing GNN minibatching. This reduces the redundant computations across
|
|
GPUs at the expense of communication.
|
|
"""
|
|
|
|
@staticmethod
|
|
def forward(
|
|
ctx,
|
|
subgraph: SampledSubgraph,
|
|
tensor: Union[torch.Tensor, Dict[str, torch.Tensor]],
|
|
):
|
|
"""Implements the forward pass."""
|
|
counts_sent = convert_to_hetero(subgraph._counts_sent)
|
|
counts_received = convert_to_hetero(subgraph._counts_received)
|
|
seed_inverse_ids = convert_to_hetero(subgraph._seed_inverse_ids)
|
|
seed_sizes = convert_to_hetero(subgraph._seed_sizes)
|
|
ctx.communication_variables = (
|
|
counts_sent,
|
|
counts_received,
|
|
seed_inverse_ids,
|
|
seed_sizes,
|
|
)
|
|
outs = {}
|
|
for ntype, typed_tensor in convert_to_hetero(tensor).items():
|
|
out = typed_tensor.new_empty(
|
|
(sum(counts_sent[ntype]),) + typed_tensor.shape[1:]
|
|
)
|
|
all_to_all(
|
|
torch.split(out, counts_sent[ntype]),
|
|
torch.split(
|
|
typed_tensor[seed_inverse_ids[ntype]],
|
|
counts_received[ntype],
|
|
),
|
|
)
|
|
outs[ntype] = out
|
|
return revert_to_homo(out)
|
|
|
|
@staticmethod
|
|
def backward(
|
|
ctx, grad_output: Union[torch.Tensor, Dict[str, torch.Tensor]]
|
|
):
|
|
"""Implements the backward pass."""
|
|
(
|
|
counts_sent,
|
|
counts_received,
|
|
seed_inverse_ids,
|
|
seed_sizes,
|
|
) = ctx.communication_variables
|
|
delattr(ctx, "communication_variables")
|
|
outs = {}
|
|
for ntype, typed_grad_output in convert_to_hetero(grad_output).items():
|
|
out = typed_grad_output.new_empty(
|
|
(sum(counts_received[ntype]),) + typed_grad_output.shape[1:]
|
|
)
|
|
all_to_all(
|
|
torch.split(out, counts_received[ntype]),
|
|
torch.split(typed_grad_output, counts_sent[ntype]),
|
|
)
|
|
i = out.new_empty(2, out.shape[0], dtype=torch.int64)
|
|
i[0] = seed_inverse_ids[ntype] # src
|
|
i[1] = torch.arange(
|
|
out.shape[0], device=typed_grad_output.device
|
|
) # dst
|
|
coo = torch.sparse_coo_tensor(
|
|
i,
|
|
torch.ones(
|
|
i.shape[1], dtype=grad_output.dtype, device=i.device
|
|
),
|
|
size=(seed_sizes[ntype], i.shape[1]),
|
|
)
|
|
outs[ntype] = torch.sparse.mm(coo, out)
|
|
return None, revert_to_homo(outs)
|
|
|
|
|
|
class CooperativeConv(torch.nn.Module):
|
|
"""Cooperative convolution operation from Cooperative Minibatching.
|
|
|
|
Implements the `all-to-all` message passing algorithm
|
|
in Cooperative Minibatching, which was initially proposed in
|
|
`Deep Graph Library PR#4337<https://github.com/dmlc/dgl/pull/4337>`__ and
|
|
was later first fully described in
|
|
`Cooperative Minibatching in Graph Neural Networks
|
|
<https://arxiv.org/abs/2310.12403>`__.
|
|
Cooperation between the GPUs eliminates duplicate work performed across the
|
|
GPUs due to the overlapping sampled k-hop neighborhoods of seed nodes when
|
|
performing GNN minibatching. This reduces the redundant computations across
|
|
GPUs at the expense of communication.
|
|
"""
|
|
|
|
def forward(
|
|
self,
|
|
subgraph: SampledSubgraph,
|
|
x: Union[torch.Tensor, Dict[str, torch.Tensor]],
|
|
):
|
|
"""Implements the forward pass."""
|
|
return CooperativeConvFunction.apply(subgraph, x)
|