Repository: pyg-team/pytorch_geometric
Pull request: Let `ImbalancedSampler` accept `torch.Tensor` as input
Hi, I modified `torch_geometric.loader.ImbalancedSampler` to accept `torch.Tensor` object, i.e., the class distribution as input. Now it is supported to use the class labels to initialize the sampler, as like:
```python
from torch_geometric.loader import NeighborLoader, ImbalancedSampler

sampler = ImbalancedSampler(data.y, input_nodes=data.train_mask)
loader = NeighborLoader(data.y, input_nodes=data.train_mask,
                        batch_size=64, num_neighbors=[-1, -1],
                        sampler=sampler, ...)    
                                sampler=sampler, ...)  
```
And the issue in #5108 should be fixed with minimal adaptation:
```python
from torch_geometric.datasets import DBLP
from torch_geometric.loader import ImbalancedSampler, NeighborLoader

dataset = DBLP(root='~/data/pygdata/dblp')
data = dataset[0]

# NOTE: input the node labels instead of the dataset `data['author']`
sampler = ImbalancedSampler(data['author'].y, input_nodes=data['author'].train_mask) 

train_loader = NeighborLoader(data, input_nodes=('author', data['author'].train_mask),
                              num_neighbors=[10, 10], 
                              sampler=sampler, batch_size=20)
```

Problem statement:
AttributeError: 'str' object has no attribute 'y' when using  torch_geometric.loader.ImbalancedSampler
### 🐛 Describe the bug

I came across a bug when using `torch_geometric.loader.ImbalancedSampler` on a Heterogeneous graph. Here is my code:
```python
from torch_geometric.datasets import DBLP
from torch_geometric.loader import ImbalancedSampler, NeighborLoader

dataset = DBLP(root='~/data/pygdata/dblp')
data = dataset[0]

sampler = ImbalancedSampler(data['author'], input_nodes=data['author'].train_mask)

train_loader = NeighborLoader(data, input_nodes=('author', data['author'].train_mask),
                              num_neighbors=[10, 10], 
                              sampler=sampler, batch_size=20)
```

The output messages are:
```
---------------------------------------------------------------------------
AttributeError                            Traceback (most recent call last)
Input In [4], in <cell line: 7>()
      4 dataset = DBLP(root='~/data/pygdata/dblp')
      5 data = dataset[0]
----> 7 sampler = ImbalancedSampler(data['author'], input_nodes=data['author'].train_mask)
      9 train_loader = NeighborLoader(data, input_nodes=('author', data['author'].train_mask),
     10                               num_neighbors=[10, 10], 
     11                               sampler=sampler, batch_size=20)

File ~/Documents/Jupyter/pytorch_geometric/torch_geometric/loader/imbalanced_sampler.py:68, in ImbalancedSampler.__init__(self, dataset, input_nodes, num_samples)
     65     assert len(dataset) == y.numel()
     67 else:
---> 68     ys = [data.y for data in dataset]
     69     if isinstance(ys[0], Tensor):
     70         y = torch.cat(ys, dim=0).view(-1)

File ~/Documents/Jupyter/pytorch_geometric/torch_geometric/loader/imbalanced_sampler.py:68, in <listcomp>(.0)
     65     assert len(dataset) == y.numel()
     67 else:
---> 68     ys = [data.y for data in dataset]
     69     if isinstance(ys[0], Tensor):
     70         y = torch.cat(ys, dim=0).view(-1)

AttributeError: 'str' object has no attribute 'y'
```
This seems attributed to the data['author'] that is not recognized as a `Data` instance.

### Environment

* PyG version: 2.0.4
* PyTorch version: 1.12
* OS: Mac OS
* Python version: 3.9
* CUDA/cuDNN version: N/A
* How you installed PyTorch and PyG (`conda`, `pip`, source): pip
* Any other relevant information (*e.g.*, version of `torch-scatter`): N/A

Hints from issue discussion:
My solution is to modify class `ImbalancedSampler`, which should accept torch.Tensor `y` as well, like this:
```python
if isinstance(dataset, torch.Tensor):
    y = y.view(-1)
    y = y[input_nodes] if input_nodes is not None else y
```

This is pretty clever solution. Do you have interest in fixing this on our end?
Sure! Could you give me some instructions? Or I just follow the contributing guidelines in https://github.com/pyg-team/pytorch_geometric/blob/master/CONTRIBUTING.md?
Yes, the contribution guide goes over some steps on how to install PyG locally, but you do not have to get too deep into it. Sending a quick PR based on a PyG fork is sufficient. We can more tests, etc as we go :) 
Thanks, I'll give it a try :)

Commit message: [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

Patch to review:
diff --git a/test/loader/test_imbalanced_sampler.py b/test/loader/test_imbalanced_sampler.py
index 9c787b8bd1be..2ce58e069efa 100644
--- a/test/loader/test_imbalanced_sampler.py
+++ b/test/loader/test_imbalanced_sampler.py
@@ -41,6 +41,7 @@ def test_neighbor_loader_with_imbalanced_sampler():
     edge_index = torch.empty((2, 0), dtype=torch.long)
     data = Data(edge_index=edge_index, y=y, num_nodes=y.size(0))
 
+    # test with `torch_geometric.data.Data` instance as input
     torch.manual_seed(12345)
     sampler = ImbalancedSampler(data)
     loader = NeighborLoader(data, batch_size=10, sampler=sampler,
@@ -55,3 +56,19 @@ def test_neighbor_loader_with_imbalanced_sampler():
 
     assert histogram.sum() == data.num_nodes
     assert prob.min() > 0.4 and prob.max() < 0.6
+
+    # test with `torch.Tensor` instance as input
+    torch.manual_seed(12345)
+    sampler = ImbalancedSampler(data.y)
+    loader = NeighborLoader(data, batch_size=10, sampler=sampler,
+                            num_neighbors=[-1])
+
+    ys: List[Tensor] = []
+    for batch in loader:
+        ys.append(batch.y)
+
+    histogram = torch.cat(ys).bincount()
+    prob = histogram / histogram.sum()
+
+    assert histogram.sum() == data.num_nodes
+    assert prob.min() > 0.4 and prob.max() < 0.6
diff --git a/torch_geometric/loader/imbalanced_sampler.py b/torch_geometric/loader/imbalanced_sampler.py
index f8850d31a34b..bac4a7de2844 100644
--- a/torch_geometric/loader/imbalanced_sampler.py
+++ b/torch_geometric/loader/imbalanced_sampler.py
@@ -33,10 +33,22 @@ class distribution.
                                 batch_size=64, num_neighbors=[-1, -1],
                                 sampler=sampler, ...)
 
+    It is also supported to accept a :class:`torch.Tensor` object,
+    which is the node label that represents the class distribution:
+
+    .. code-block:: python
+
+        from torch_geometric.loader import NeighborLoader, ImbalancedSampler
+
+        sampler = ImbalancedSampler(data.y, input_nodes=data.train_mask)
+        loader = NeighborLoader(data.y, input_nodes=data.train_mask,
+                                batch_size=64, num_neighbors=[-1, -1],
+                                sampler=sampler, ...)
+
     Args:
-        dataset (Dataset or Data): The dataset from which to sample the data,
-            either given as a :class:`~torch_geometric.data.Dataset` or
-            :class:`~torch_geometric.data.Data` object.
+        dataset (Dataset, Data, or Tensor): The dataset or class distribution
+            from which to sample the data, given as a :class:`~torch_geometric.data.Dataset`,
+            :class:`~torch_geometric.data.Data` or :class:`torch.Tensor` object.
         input_nodes (Tensor, optional): The indices of nodes that are used by
             the corresponding loader, *e.g.*, by
             :class:`~torch_geometric.loader.NeighborLoader`.
@@ -50,7 +62,7 @@ class distribution.
     """
     def __init__(
         self,
-        dataset: Union[Data, Dataset, List[Data]],
+        dataset: Union[Data, Tensor, Dataset, List[Data]],
         input_nodes: Optional[Tensor] = None,
         num_samples: Optional[int] = None,
     ):
@@ -60,6 +72,10 @@ def __init__(
             assert dataset.num_nodes == y.numel()
             y = y[input_nodes] if input_nodes is not None else y
 
+        elif isinstance(dataset, Tensor):
+            y = dataset.view(-1)
+            y = y[input_nodes] if input_nodes is not None else y
+
         elif isinstance(dataset, InMemoryDataset):
             y = dataset.data.y.view(-1)
             assert len(dataset) == y.numel()


Code review comments:
diff_hunk: @@ -33,10 +33,22 @@ class distribution.
                                 batch_size=64, num_neighbors=[-1, -1],
                                 sampler=sampler, ...)
 
+    It is also supported to accept a :class:`torch.Tensor` object,
+    which is the node label that represents the class distribution:
path: torch_geometric/loader/imbalanced_sampler.py
text: @user1:
```suggestion
    which has the node label distribution that we wish to sample with:
```
Is this accurate? (No need to accept this change, but the current description was confusing for me)

I would have thought you'd directly want to provide weights in this case.