NeighborSampler samples different nodes with disjoint=True and False

[Barcavin]

🐛 Describe the bug

import argparse

from torch_geometric.datasets import FakeDataset, FakeHeteroDataset
from torch_geometric.loader import NeighborLoader
from torch_geometric.seed import seed_everything

parser = argparse.ArgumentParser()
parser.add_argument('--disjoint', type=int, default=0)
args = parser.parse_args()
seed_everything(0)

data = FakeHeteroDataset(1, avg_degree=600)[0]

print(data['v0'].x.mean())
loader = NeighborLoader(
            data,
            # Sample 30 neighbors for each node for 2 iterations
            num_neighbors=[30] * 2,
            # Use a batch size of 128 for sampling training nodes
            batch_size=128,
            input_nodes=('v0',[1]),
            disjoint=args.disjoint,
)
print(next(iter(loader))['v0'].n_id.unique().sum())

Run in the terminal:

python disjoint.py --disjoint=1

tensor(-0.0022)
tensor(72436)

python disjoint.py --disjoint=0

tensor(-0.0022)
tensor(71816)

On the same datasets, it will sample two different node sets. I only see this happens on heterogeneous data but not on homogeneous data.

Versions

PyTorch version: 2.0.1
Is debug build: False
CUDA used to build PyTorch: 11.8
ROCM used to build PyTorch: N/A

OS: Ubuntu 22.04.3 LTS (x86_64)
GCC version: (Ubuntu 11.4.0-1ubuntu1~22.04) 11.4.0
Clang version: Could not collect
CMake version: Could not collect
Libc version: glibc-2.35

Python version: 3.11.9 | packaged by conda-forge | (main, Apr 19 2024, 18:36:13) [GCC 12.3.0] (64-bit runtime)
Python platform: Linux-4.14.343-260.564.amzn2.x86_64-x86_64-with-glibc2.35
Is CUDA available: True
CUDA runtime version: Could not collect
CUDA_MODULE_LOADING set to: LAZY
GPU models and configuration: GPU 0: NVIDIA A10G
Nvidia driver version: 535.129.03
cuDNN version: Could not collect
HIP runtime version: N/A
MIOpen runtime version: N/A
Is XNNPACK available: True

CPU:
Architecture: x86_64
CPU op-mode(s): 32-bit, 64-bit
Address sizes: 48 bits physical, 48 bits virtual
Byte Order: Little Endian
CPU(s): 32
On-line CPU(s) list: 0-31
Vendor ID: AuthenticAMD
Model name: AMD EPYC 7R32
CPU family: 23
Model: 49
Thread(s) per core: 2
Core(s) per socket: 16
Socket(s): 1
Stepping: 0
BogoMIPS: 5599.99
Flags: fpu vme de pse tsc msr pae mce cx8 apic sep mtrr pge mca cmov pat pse36 clflush mmx fxsr sse sse2 ht syscall nx mmxext fxsr_opt pdpe1gb rdtscp lm constant_tsc rep_good nopl xtopology nonstop_tsc cpuid extd_apicid aperfmperf tsc_known_freq pni pclmulqdq ssse3 fma cx16 sse4_1 sse4_2 movbe popcnt aes xsave avx f16c rdrand hypervisor lahf_lm cmp_legacy cr8_legacy abm sse4a misalignsse 3dnowprefetch topoext ssbd ibrs ibpb stibp vmmcall fsgsbase bmi1 avx2 smep bmi2 rdseed adx smap clflushopt clwb sha_ni xsaveopt xsavec xgetbv1 clzero xsaveerptr arat npt nrip_save rdpid
Hypervisor vendor: KVM
Virtualization type: full
L1d cache: 512 KiB (16 instances)
L1i cache: 512 KiB (16 instances)
L2 cache: 8 MiB (16 instances)
L3 cache: 64 MiB (4 instances)
NUMA node(s): 1
NUMA node0 CPU(s): 0-31
Vulnerability Gather data sampling: Not affected
Vulnerability Itlb multihit: Not affected
Vulnerability L1tf: Not affected
Vulnerability Mds: Not affected
Vulnerability Meltdown: Not affected
Vulnerability Mmio stale data: Not affected
Vulnerability Retbleed: Vulnerable, RAS-Poisoning: Vulnerable
Vulnerability Spec store bypass: Mitigation; Speculative Store Bypass disabled via prctl and seccomp
Vulnerability Spectre v1: Mitigation; usercopy/swapgs barriers and __user pointer sanitization
Vulnerability Spectre v2: Mitigation; Retpolines, IBPB conditional, IBRS_FW, STIBP conditional, RSB filling, PBRSB-eIBRS Not affected
Vulnerability Srbds: Not affected
Vulnerability Tsx async abort: Not affected

Versions of relevant libraries:
[pip3] numpy==1.26.4
[pip3] pytorch_frame==0.2.2
[pip3] torch==2.0.1
[pip3] torch_geometric==2.5.2
[pip3] torch-hd==5.5.0
[pip3] torch-scatter==2.1.2
[pip3] torch-sparse==0.6.18
[pip3] torchaudio==2.0.2
[pip3] torchvision==0.15.2
[pip3] triton==2.0.0
[conda] blas 2.121 mkl conda-forge
[conda] blas-devel 3.9.0 21_linux64_mkl conda-forge
[conda] ffmpeg 4.3 hf484d3e_0 pytorch
[conda] libblas 3.9.0 21_linux64_mkl conda-forge
[conda] libcblas 3.9.0 21_linux64_mkl conda-forge
[conda] liblapack 3.9.0 21_linux64_mkl conda-forge
[conda] liblapacke 3.9.0 21_linux64_mkl conda-forge
[conda] mkl 2024.0.0 ha957f24_49657 conda-forge
[conda] mkl-devel 2024.0.0 ha770c72_49657 conda-forge
[conda] mkl-include 2024.0.0 ha957f24_49657 conda-forge
[conda] numpy 1.26.4 py311h64a7726_0 conda-forge
[conda] pyg 2.5.2 py311_torch_2.0.0_cu118 pyg
[conda] pytorch 2.0.1 py3.11_cuda11.8_cudnn8.7.0_0 pytorch
[conda] pytorch-cuda 11.8 h7e8668a_5 pytorch
[conda] pytorch-frame 0.2.2 pypi_0 pypi
[conda] pytorch-mutex 1.0 cuda pytorch
[conda] pytorch-scatter 2.1.2 py311_torch_2.0.0_cu118 pyg
[conda] pytorch-sparse 0.6.18 py311_torch_2.0.0_cu118 pyg
[conda] torch-hd 5.5.0 pypi_0 pypi
[conda] torchaudio 2.0.2 py311_cu118 pytorch
[conda] torchtriton 2.0.0 py311 pytorch
[conda] torchvision 0.15.2 py311_cu118 pytorch

[rusty1s]

I would say that is expected. Although you fix the random seed, sampling is performed differently for disjoint=True/False. You can see that the same number of nodes are sampled if you change to num_neighbors=[-1] * 2.

[Barcavin]

Thanks for your reply. However, the sampling seems to be more consistent on homogeneous data. Is that also expected somehow on heterogeneous data?

[rusty1s]

What do you mean by "more consistent"?

[Barcavin]

If we sample on homogeneous graph, the sampled data for either disjoint=True or False will be the same.

To reproduce, running python disjoint.py --disjoint=0 --heter=0 and python disjoint.py --disjoint=1 --heter=0 will give the same result.

disjoint.py:

import argparse

from torch_geometric.datasets import FakeDataset, FakeHeteroDataset
from torch_geometric.loader import NeighborLoader
from torch_geometric.seed import seed_everything

parser = argparse.ArgumentParser()
parser.add_argument('--disjoint', type=int, default=0)
parser.add_argument('--heter', type=int, default=1)
args = parser.parse_args()
seed_everything(0)

if args.heter:
    data = FakeHeteroDataset(1, avg_degree=600)[0]

    print(data['v0'].x.mean())
    loader = NeighborLoader(
                data,
                # Sample 30 neighbors for each node for 2 iterations
                num_neighbors=[30] * 2,
                # Use a batch size of 128 for sampling training nodes
                batch_size=128,
                input_nodes=('v0',[1]),
                disjoint=args.disjoint,
    )
    print(next(iter(loader))['v0'].n_id.unique().sum())
else:
    data = FakeDataset(1, avg_degree=600)[0]

    print(data.x.mean())
    loader = NeighborLoader(
                data,
                # Sample 30 neighbors for each node for 2 iterations
                num_neighbors=[30] * 2,
                # Use a batch size of 128 for sampling training nodes
                batch_size=128,
                input_nodes=[1],
                disjoint=args.disjoint,
    )
    print(next(iter(loader)).n_id.unique().sum())

[rusty1s]

Thanks. I don't necessarily think this is a problem, but I will try to find some time to look into why this is the case.