PyTorch Geometric: Accessing batch in a custom "In Memory" dataset fails

dokondr: <p>When trying to get batch in dataloader I get a key error from <code>torch_geometric\data\storage.py</code>. Please see below:</p>
<p>import torch<br>
from torch_geometric.data import Data, InMemoryDataset<br>
import matplotlib.pyplot as plt<br>
import networkx as nx</p>
<h1>
<a name="define-the-edge-index-and-features-for-the-first-graph-1" class="anchor" href="#define-the-edge-index-and-features-for-the-first-graph-1"></a>Define the edge index and features for the first graph</h1>
<p>x1 = torch.tensor([[1, 2], [3, 4]], dtype=torch.float)<br>
edge_index1 = torch.tensor([[0, 1, 1, 0], [1, 0, 0, 1]], dtype=torch.long)</p>
<h1>
<a name="define-the-edge-index-and-features-for-the-second-graph-2" class="anchor" href="#define-the-edge-index-and-features-for-the-second-graph-2"></a>Define the edge index and features for the second graph</h1>
<p>x2 = torch.tensor([[2, 3], [1, 4]], dtype=torch.float)<br>
edge_index2 = torch.tensor([[0, 1, 1, 0], [1, 0, 0, 1]], dtype=torch.long)</p>
<h1>
<a name="create-a-list-of-pytorch-geometric-data-objects-to-represent-the-graphs-3" class="anchor" href="#create-a-list-of-pytorch-geometric-data-objects-to-represent-the-graphs-3"></a>Create a list of PyTorch Geometric Data objects to represent the graphs</h1>
<p>data_list = [Data(x=x1, edge_index=edge_index1),<br>
Data(x=x2, edge_index=edge_index2)]</p>
<h1>
<a name="define-the-custom-dataset-class-4" class="anchor" href="#define-the-custom-dataset-class-4"></a>Define the custom dataset class</h1>
<p>class CustomDataset(InMemoryDataset):<br>
def <strong>init</strong>(self, root, transform=None, pre_transform=None):<br>
super(CustomDataset, self).<strong>init</strong>(root, transform, pre_transform)<br>
self.data, self.slices = self.collate(data_list)</p>
<pre><code>  @property
  def raw_file_names(self):
      return []

  @property
  def processed_file_names(self):
      return []

  def download(self):
      pass

  def process(self):
      pass

  def __len__(self):
      return len(self.data)

  def get(self, idx):
      # Return the Data object at the specified index
      return self.data[idx]
</code></pre>
<h1>
<a name="initialize-the-custom-dataset-5" class="anchor" href="#initialize-the-custom-dataset-5"></a>Initialize the custom dataset</h1>
<p>dataset = CustomDataset(root=‘./data’)</p>
<h1>
<a name="apply-transformations-to-the-data-if-desired-6" class="anchor" href="#apply-transformations-to-the-data-if-desired-6"></a>Apply transformations to the data if desired</h1>
<p>if dataset.transform is not None:<br>
dataset.transform = T.Compose([T.RandomRotate(30, resample=False),<br>
T.RandomTranslate(0.1)])</p>
<h1>
<a name="put-the-dataset-in-a-data-loader-7" class="anchor" href="#put-the-dataset-in-a-data-loader-7"></a>Put the dataset in a data loader</h1>
<p>dataloader = torch.utils.data.DataLoader(dataset, batch_size=2, shuffle=True)</p>
<h1>
<a name="loop-over-the-dataloader-8" class="anchor" href="#loop-over-the-dataloader-8"></a>Loop over the dataloader</h1>
<p>for batch in dataloader:<br>
# Access the data for each graph in the batch<br>
# x, edge_index, y = batch.x, batch.edge_index, batch.y<br>
x, edge_index = batch.x, batch.edge_index</p>
<pre><code>  # Do something with the data
  print(x, edge_index, y)
</code></pre>
<p>######## Error #######<br>
KeyError                                  Traceback (most recent call last)<br>
Cell In[4], line 2<br>
1 # Loop over the dataloader<br>
----&gt; 2 for batch in dataloader:<br>
3     # Access the data for each graph in the batch<br>
4     # x, edge_index, y = batch.x, batch.edge_index, batch.y<br>
5     x, edge_index = batch.x, batch.edge_index<br>
7     # Do something with the data</p>
<p>File ~\Anaconda3\envs\pyg\lib\site-packages\torch\utils\data\dataloader.py:681, in _BaseDataLoaderIter.<strong>next</strong>(self)<br>
678 if self._sampler_iter is None:<br>
679     # TODO(<a href="https://github.com/pytorch/pytorch/issues/76750" class="inline-onebox" rel="noopener nofollow ugc">Bug in dataloader iterator found by mypy · Issue #76750 · pytorch/pytorch · GitHub</a>)<br>
680     self._reset()  # type: ignore[call-arg]<br>
 → 681 data = self._next_data()<br>
682 self._num_yielded += 1<br>
683 if self._dataset_kind == _DatasetKind.Iterable and <br>
684         self._IterableDataset_len_called is not None and <br>
685         self._num_yielded &gt; self._IterableDataset_len_called:</p>
<p>File ~\Anaconda3\envs\pyg\lib\site-packages\torch\utils\data\dataloader.py:721, in _SingleProcessDataLoaderIter._next_data(self)<br>
719 def _next_data(self):<br>
720     index = self._next_index()  # may raise StopIteration<br>
 → 721     data = self._dataset_fetcher.fetch(index)  # may raise StopIteration<br>
722     if self._pin_memory:<br>
723         data = _utils.pin_memory.pin_memory(data, self._pin_memory_device)</p>
<p>File ~\Anaconda3\envs\pyg\lib\site-packages\torch\utils\data_utils\fetch.py:49, in _MapDatasetFetcher.fetch(self, possibly_batched_index)<br>
47 def fetch(self, possibly_batched_index):<br>
48     if self.auto_collation:<br>
—&gt; 49         data = [self.dataset[idx] for idx in possibly_batched_index]<br>
50     else:<br>
51         data = self.dataset[possibly_batched_index]</p>
<p>File ~\Anaconda3\envs\pyg\lib\site-packages\torch\utils\data_utils\fetch.py:49, in (.0)<br>
47 def fetch(self, possibly_batched_index):<br>
48     if self.auto_collation:<br>
—&gt; 49         data = [self.dataset[idx] for idx in possibly_batched_index]<br>
50     else:<br>
51         data = self.dataset[possibly_batched_index]</p>
<p>File ~\Anaconda3\envs\pyg\lib\site-packages\torch_geometric\data\dataset.py:197, in Dataset.<strong>getitem</strong>(self, idx)<br>
187 r"““In case :obj:<code>idx</code> is of type integer, will return the data object<br>
188 at index :obj:<code>idx</code> (and transforms it in case :obj:<code>transform</code> is<br>
189 present).<br>
190 In case :obj:<code>idx</code> is a slicing object, <em>e.g.</em>, :obj:<code>[2:5]</code>, a list, a<br>
191 tuple, or a :obj:<code>torch.Tensor</code> or :obj:<code>np.ndarray</code> of type long or<br>
192 bool, will return a subset of the dataset at the specified indices.””"<br>
193 if (isinstance(idx, (int, np.integer))<br>
194         or (isinstance(idx, Tensor) and idx.dim() == 0)<br>
195         or (isinstance(idx, np.ndarray) and np.isscalar(idx))):<br>
 → 197     data = self.get(self.indices()[idx])<br>
198     data = data if self.transform is None else self.transform(data)<br>
199     return data</p>
<p>Cell In[3], line 43, in CustomDataset.get(self, idx)<br>
41 def get(self, idx):<br>
42     # Return the Data object at the specified index<br>
—&gt; 43     return self.data[idx]</p>
<p>File ~\Anaconda3\envs\pyg\lib\site-packages\torch_geometric\data\data.py:444, in Data.<strong>getitem</strong>(self, key)<br>
443 def <strong>getitem</strong>(self, key: str) → Any:<br>
 → 444     return self._store[key]</p>
<p>File ~\Anaconda3\envs\pyg\lib\site-packages\torch_geometric\data\storage.py:81, in BaseStorage.<strong>getitem</strong>(self, key)<br>
80 def <strong>getitem</strong>(self, key: str) → Any:<br>
—&gt; 81     return self._mapping[key]</p>
<p>KeyError: 0</p>