mini-batching issue - can't iterate over batches

[YoavShamir5]

Hi,

I have a strange issue - I am generating a DataLoader, but then I am unable to iterate over the batches.

Here is the code:

from torch_geometric.data import Dataset, download_url, Data
import os 
from torch_geometric.loader import DataLoader
import torch

class WarheadDataset(Dataset):
    def __init__(self, root, transform=None, pre_transform=None, pre_filter=None):
        # self.processed_dir = "processed_data_dir_dock_2"
        super().__init__(root, transform, pre_transform, pre_filter)

    @property
    def processed_dir(self) -> str:
        return os.path.join(self.root, 'data_files_try')

    @property
    def processed_file_names(self):
        return [filename for filename in os.listdir(self.processed_dir)]

    def len(self):
        return len(self.processed_file_names)

    def get(self, idx):
        print("idx: ", idx)
        data = torch.load(os.path.join(self.processed_dir, f'data_{idx}.pt'))
        print(data)

        return data

dataset_root_dir = r"C:\Users\SYOAV\Desktop\PDBbind"

dataset = WarheadDataset(root = dataset_root_dir)

dataset_loader = DataLoader(dataset, batch_size = 3)

for batch in dataset_loader:
    print(batch)

here is the output:

idx:  0
Data(x=[29], pos=[29], lig_bit=[29])
idx:  1
Data(x=[14], pos=[14], lig_bit=[14])
idx:  2
Data(x=[36], pos=[36], lig_bit=[36])

here is the error message:

---------------------------------------------------------------------------
IndexError                                Traceback (most recent call last)
Cell In[29], line 45
     42 dataset_loader = DataLoader(dataset, batch_size = 3)
     43 load_iter = iter(dataset_loader)
---> 45 for batch in dataset_loader:
     46     print(batch)

File [c:\Users\SYOAV\anaconda3\envs\envl\lib\site-packages\torch\utils\data\dataloader.py:633](file:///C:/Users/SYOAV/anaconda3/envs/envl/lib/site-packages/torch/utils/data/dataloader.py:633), in _BaseDataLoaderIter.__next__(self)
    630 if self._sampler_iter is None:
    631     # TODO(https://github.com/pytorch/pytorch/issues/76750)
    632     self._reset()  # type: ignore[call-arg]
--> 633 data = self._next_data()
    634 self._num_yielded += 1
    635 if self._dataset_kind == _DatasetKind.Iterable and \
    636         self._IterableDataset_len_called is not None and \
    637         self._num_yielded > self._IterableDataset_len_called:

File [c:\Users\SYOAV\anaconda3\envs\envl\lib\site-packages\torch\utils\data\dataloader.py:677](file:///C:/Users/SYOAV/anaconda3/envs/envl/lib/site-packages/torch/utils/data/dataloader.py:677), in _SingleProcessDataLoaderIter._next_data(self)
    675 def _next_data(self):
    676     index = self._next_index()  # may raise StopIteration
--> 677     data = self._dataset_fetcher.fetch(index)  # may raise StopIteration
    678     if self._pin_memory:
    679         data = _utils.pin_memory.pin_memory(data, self._pin_memory_device)

File [c:\Users\SYOAV\anaconda3\envs\envl\lib\site-packages\torch\utils\data\_utils\fetch.py:54](file:///C:/Users/SYOAV/anaconda3/envs/envl/lib/site-packages/torch/utils/data/_utils/fetch.py:54), in _MapDatasetFetcher.fetch(self, possibly_batched_index)
     52 else:
     53     data = self.dataset[possibly_batched_index]
---> 54 return self.collate_fn(data)

File [c:\Users\SYOAV\anaconda3\envs\envl\lib\site-packages\torch_geometric\loader\dataloader.py:20](file:///C:/Users/SYOAV/anaconda3/envs/envl/lib/site-packages/torch_geometric/loader/dataloader.py:20), in Collater.__call__(self, batch)
     18 elem = batch[0]
     19 if isinstance(elem, BaseData):
---> 20     return Batch.from_data_list(batch, self.follow_batch,
     21                                 self.exclude_keys)
     22 elif isinstance(elem, torch.Tensor):
     23     return default_collate(batch)

File [c:\Users\SYOAV\anaconda3\envs\envl\lib\site-packages\torch_geometric\data\batch.py:76](file:///C:/Users/SYOAV/anaconda3/envs/envl/lib/site-packages/torch_geometric/data/batch.py:76), in Batch.from_data_list(cls, data_list, follow_batch, exclude_keys)
     64 @classmethod
     65 def from_data_list(cls, data_list: List[BaseData],
     66                    follow_batch: Optional[List[str]] = None,
     67                    exclude_keys: Optional[List[str]] = None):
     68     r"""Constructs a :class:`~torch_geometric.data.Batch` object from a
     69     Python list of :class:`~torch_geometric.data.Data` or
     70     :class:`~torch_geometric.data.HeteroData` objects.
   (...)
     73     :obj:`follow_batch`.
     74     Will exclude any keys given in :obj:`exclude_keys`."""
---> 76     batch, slice_dict, inc_dict = collate(
     77         cls,
     78         data_list=data_list,
     79         increment=True,
     80         add_batch=not isinstance(data_list[0], Batch),
     81         follow_batch=follow_batch,
     82         exclude_keys=exclude_keys,
     83     )
     85     batch._num_graphs = len(data_list)
     86     batch._slice_dict = slice_dict

File [c:\Users\SYOAV\anaconda3\envs\envl\lib\site-packages\torch_geometric\data\collate.py:84](file:///C:/Users/SYOAV/anaconda3/envs/envl/lib/site-packages/torch_geometric/data/collate.py:84), in collate(cls, data_list, increment, add_batch, follow_batch, exclude_keys)
     81     continue
     83 # Collate attributes into a unified representation:
---> 84 value, slices, incs = _collate(attr, values, data_list, stores,
     85                                increment)
     87 if isinstance(value, Tensor) and value.is_cuda:
     88     device = value.device

File [c:\Users\SYOAV\anaconda3\envs\envl\lib\site-packages\torch_geometric\data\collate.py:194](file:///C:/Users/SYOAV/anaconda3/envs/envl/lib/site-packages/torch_geometric/data/collate.py:194), in _collate(key, values, data_list, stores, increment)
    192 value_list, slice_list, inc_list = [], [], []
    193 for i in range(len(elem)):
--> 194     value, slices, incs = _collate(key, [v[i] for v in values],
    195                                    data_list, stores, increment)
    196     value_list.append(value)
    197     slice_list.append(slices)

File [c:\Users\SYOAV\anaconda3\envs\envl\lib\site-packages\torch_geometric\data\collate.py:194](file:///C:/Users/SYOAV/anaconda3/envs/envl/lib/site-packages/torch_geometric/data/collate.py:194), in (.0)
    192 value_list, slice_list, inc_list = [], [], []
    193 for i in range(len(elem)):
--> 194     value, slices, incs = _collate(key, [v[i] for v in values],
    195                                    data_list, stores, increment)
    196     value_list.append(value)
    197     slice_list.append(slices)

IndexError: list index out of range

The same code worked for a different dataset, in both cases I have previously generated torch_geometric Data instances and saved them in individual files, I am not sure what specifically in this case is the issue. Thanks!

[rusty1s]

Do you mind sharing a few examples of your data objects? I am not totally sure what‘s wrong here :(