wehub-resource-sync
58 lines
2.1 KiB
Python
58 lines
2.1 KiB
Python
"""An abstract base class for cugraph-ops nn module."""
|
|
import torch
|
|
from torch import nn
|
|
|
|
|
|
class CuGraphBaseConv(nn.Module):
|
|
r"""An abstract base class for cugraph-ops nn module."""
|
|
|
|
def __init__(self):
|
|
super().__init__()
|
|
self._cached_offsets_fg = None
|
|
|
|
def reset_parameters(self):
|
|
r"""Resets all learnable parameters of the module."""
|
|
raise NotImplementedError
|
|
|
|
def forward(self, *args):
|
|
r"""Runs the forward pass of the module."""
|
|
raise NotImplementedError
|
|
|
|
def pad_offsets(self, offsets: torch.Tensor, size: int) -> torch.Tensor:
|
|
r"""Pad zero-in-degree nodes to the end of offsets to reach size.
|
|
|
|
cugraph-ops often provides two variants of aggregation functions for a
|
|
specific model: one intended for sampled-graph use cases, one for
|
|
full-graph ones. The former is in general more performant, however, it
|
|
only works when the sample size (the max of in-degrees) is small (<200),
|
|
due to the limit of GPU shared memory. For graphs with a larger max
|
|
in-degree, we need to fall back to the full-graph option, which requires
|
|
to convert a DGL block to a full graph. With the csc-representation,
|
|
this is equivalent to pad zero-in-degree nodes to the end of the offsets
|
|
array (also called indptr or colptr).
|
|
|
|
Parameters
|
|
----------
|
|
offsets :
|
|
The (monotonically increasing) index pointer array in a CSC-format
|
|
graph.
|
|
size : int
|
|
The length of offsets after padding.
|
|
|
|
Returns
|
|
-------
|
|
torch.Tensor
|
|
The augmented offsets array.
|
|
"""
|
|
if self._cached_offsets_fg is None:
|
|
self._cached_offsets_fg = torch.empty(
|
|
size, dtype=offsets.dtype, device=offsets.device
|
|
)
|
|
elif self._cached_offsets_fg.numel() < size:
|
|
self._cached_offsets_fg.resize_(size)
|
|
|
|
self._cached_offsets_fg[: offsets.numel()] = offsets
|
|
self._cached_offsets_fg[offsets.numel() : size] = offsets[-1]
|
|
|
|
return self._cached_offsets_fg[:size]
|