How to use NeighborLoader to sample all nodes in the graph (also isolated ones) with remote backends?

[denadai2]

Hi,
I'm trying to use NeighborLoader to sample all nodes in the graph (also isolated ones) with remote backends. How can I do that? I am able to sample just the nodes that have at least one edge. I saw the source code and I see that here pyg computes the number of nodes through edges (

pytorch_geometric/torch_geometric/sampler/neighbor_sampler.py

Line 230 in a8671b1

self.num_nodes = {

). Is it correct?

Is it normal this test fails?

import pytest
import torch

from torch_geometric.data import HeteroData
from torch_geometric.data.remote_backend_utils import num_nodes, size
from torch_geometric.testing import (
    MyFeatureStore,
    MyGraphStore,
    get_random_edge_index,
)

@pytest.mark.parametrize('FeatureStore', [MyFeatureStore, HeteroData])
@pytest.mark.parametrize('GraphStore', [MyGraphStore, HeteroData])
def test_mything(FeatureStore, GraphStore):
    feature_store = FeatureStore()
    graph_store = GraphStore()

    x = torch.arange(100)
    feature_store.put_tensor(x, group_name='x', attr_name='x', index=None)

    # Infer num nodes and size from edges:
    xy = get_random_edge_index(90, 50, 20)
    graph_store.put_edge_index(xy, edge_type=('x', 'to', 'y'), layout='coo',
                               size=(90, 50))

    assert size(feature_store, graph_store, 'x') == 100

thxx

[rusty1s]

Why don't you just set size=(100, 50)?

[rusty1s]

Yes, if they are part of the seed nodes/input nodes. However, since they are not reachable, these nodes will not get sampled from other seed nodes.

[denadai2]

so what am I doing wrong here? :( sorry for the question

import pytest
import torch

from torch_geometric.data import HeteroData
from torch_geometric.data.remote_backend_utils import num_nodes, size
from torch_geometric.testing import (
    MyFeatureStore,
    MyGraphStore,
    get_random_edge_index,
)
from torch_geometric.loader import NeighborLoader


@pytest.mark.parametrize('FeatureStore', [MyFeatureStore, HeteroData])
@pytest.mark.parametrize('GraphStore', [MyGraphStore, HeteroData])
def test_num_nodes_size(FeatureStore, GraphStore):
    feature_store = FeatureStore()
    graph_store = GraphStore()

    # Infer num nodes from features:
    x = torch.arange(100)
    feature_store.put_tensor(x, group_name='x', attr_name='x', index=None)
    assert num_nodes(feature_store, graph_store, 'x') == 100

    # Infer num nodes and size from edges:
    xy = get_random_edge_index(90, 50, 20)
    graph_store.put_edge_index(xy, edge_type=('x', 'to', 'y'), layout='coo',
                               size=(90, 50))

    loader = NeighborLoader(
            (feature_store, graph_store),
            num_neighbors={
                ('x', 'to', 'y'): [-1],
            },
            input_nodes='x',
            batch_size=1,
    )
    assert len(loader) == 100

the last assert fails. I would expect to loop through all 100 nodes. One way to fix it would be to add self loops, but I do not wanna do that

[denadai2]

I might have fixed the bug I think we have #9005

[rusty1s]

You just need to do

graph_store.put_edge_index(xy, edge_type=('x', 'to', 'y'), layout='coo',
                               size=(100, 50))

in order to fix this :)

[denadai2]

oh wow, that's true from the code. So is my PR an improvement or not?

btw: it seems your code works only with coo and not csc. Is it normal? (I need csc at the moment)

[rusty1s]

IMO, we don't necessarily need your PR since it is safer to ensure that edge_index has the correct shape. Can you post an example where CSC would fail?

[denadai2]

I see thx.

graph_store.put_edge_index(xy, edge_type=('x', 'to', 'y'), layout='csc',
                               size=(90, 50))

I just replaced this so I might be wrong

def test_neighborloader_featurestore(FeatureStore, GraphStore):
    feature_store = FeatureStore()
    graph_store = GraphStore()

    # Infer num nodes from features:
    x = torch.arange(100)
    feature_store.put_tensor(x, group_name='x', attr_name='x', index=None)

    # Infer num nodes and size from edges:
    xy = get_random_edge_index(90, 50, 20)
    graph_store.put_edge_index(xy, edge_type=('x', 'to', 'y'), layout='csc',
                               size=(90, 50))

    loader = NeighborLoader(
        (feature_store, graph_store),
        num_neighbors={
            ('x', 'to', 'y'): [-1],
        },
        input_nodes='x',
        batch_size=1,
    )
    assert len(loader) == 100

    loader = LinkNeighborLoader((feature_store, graph_store),
                                num_neighbors={
                                    ('x', 'to', 'y'): [-1],
                                },
                                edge_label_index=("x", "to", "y"),)
    assert len(loader) == 20

I have

        assert row is not None or rowptr is not None
        assert col is not None
        assert col.dtype == torch.long
        assert col.dim() == 1
        col = col.contiguous()

        M: int = 0
        if sparse_sizes is None or sparse_sizes[0] is None:
            if rowptr is not None:
                M = rowptr.numel() - 1
            elif row is not None and row.numel() > 0:
                M = int(row.max()) + 1
        else:
            _M = sparse_sizes[0]
            assert _M is not None
            M = _M
            if rowptr is not None:
>               assert rowptr.numel() - 1 == M
E               AssertionError

../../../miniconda3/envs/gnn/lib/python3.10/site-packages/torch_sparse/storage.py:67: AssertionError

[rusty1s]

Well, your input is not a CSC tensor but a COO tensor, so this is expected to fail.