[Neighborhood Sampler] Sample twice with same input nodes(root node)

[Junseok0207]

I want to sample different neighbors twice with same input(root) nodes.
It is too hard to customize because of the complex source code for me...
Can you help me?

I change the transform_fn on NeighborLoader and call on NeighborSampler like this
but I think that same neighbors are sampled...

def transform_fn(self, out: Any) -> Union[Data, HeteroData]:
        
        view1, view2 = out
        node_dict1, row_dict1, col_dict1, edge_dict1, batch_size = view1
        data1 = filter_hetero_data(self.data, node_dict1, row_dict1, col_dict1,
                                    edge_dict1,
                                    self.neighbor_sampler.perm_dict)

        data1[self.neighbor_sampler.input_node_type].batch_size = batch_size

        node_dict2, row_dict2, col_dict2, edge_dict2, batch_size = view2
        data2 = filter_hetero_data(self.data, node_dict2, row_dict2, col_dict2,
                                    edge_dict2,
                                    self.neighbor_sampler.perm_dict)
        data2[self.neighbor_sampler.input_node_type].batch_size = batch_size

        return data1, data2

def __call__(self, indices: List[int]):
        if not isinstance(indices, Tensor):
            index = torch.tensor(indices)
        assert index.dtype == torch.int64

        sample_fn = torch.ops.torch_sparse.hetero_neighbor_sample
        node_dict1, row_dict1, col_dict1, edge_dict1 = sample_fn(
            self.node_types,
            self.edge_types,
            self.colptr_dict,
            self.row_dict,
            {self.input_node_type: index},
            self.num_neighbors,
            self.num_hops,
            self.replace,
            self.directed,
        )

        view1 = (node_dict1, row_dict1, col_dict1, edge_dict1, index.numel())

        sample_fn2 = torch.ops.torch_sparse.hetero_neighbor_sample
        node_dict2, row_dict2, col_dict2, edge_dict2 = sample_fn2(
            self.node_types,
            self.edge_types,
            self.colptr_dict,
            self.row_dict,
            {self.input_node_type: index},
            self.num_neighbors,
            self.num_hops,
            self.replace,
            self.directed,
        )

        view2 = (node_dict2, row_dict2, col_dict2, edge_dict2, index.numel())

        return view1, view2

[rusty1s]

How about doing this on a Python-level?

sampler = torch.utils.data.RandomSampler(range(num_nodes))
loader1 = NeighborLoader(data, sampler=sampler, ...)
loader2 = NeighborLoader(data, sampler=sampler, ...)

for data1, data2 in zip(loader1, loader2):
    pass

[Junseok0207]

I try this and it outputs two different outputs(data1, data2)
However, when I check with below codes

sampled_id1 = data1['paper'].n_id[:batch_size]
sampled_id2 = data2['paper'].n_id[:batch_size]

My output is

tensor([ 157, 2319, 3692,  ..., 1384, 1668, 1561])
tensor([1615, 4121, 2584,  ..., 3818,  851, 3982])

respectively, so I think that this two outputs do not have same input(root) nodes...

[rusty1s]

Oh, you are right. Instead, one needs to pass in a generator with the same manual seed to it, see https://pytorch.org/docs/stable/generated/torch.Generator.html#torch.Generator.

[Junseok0207]

It works well! thank you for your kind help!!!