How to store Data objects in Numpy array without the "conversion" of the Data object to Numpy array?

[jagandecapri]

The following code:

import torch
from torch_geometric.data import Data

x = torch.randn(10, 1)
edge_index = torch.randint(10, size=(2, 20))

a = Data(x=x, edge_index=edge_index)
b = np.array([a], dtype=object)
print(b)
print(f"Type: {type(b[0])}")

Gives the following output:

[[['x' tensor([[-1.0290],
               [ 0.7104],
               [-0.1062],
               [-0.9715],
               [ 0.0865],
               [-0.3081],
               [-0.0390],
               [ 0.6114],
               [-0.6279],
               [ 1.4209]])]
  ['edge_index'
   tensor([[9, 6, 2, 2, 9, 8, 9, 5, 4, 2, 9, 0, 4, 0, 5, 2, 8, 8, 1, 2],
           [2, 5, 4, 6, 7, 2, 1, 3, 7, 6, 9, 2, 5, 8, 3, 7, 0, 3, 2, 6]])]]]
Type: <class 'numpy.ndarray'>

The Data object seems to be "converted" to a Numpy array too. I would like to store the Data object as it is in the Numpy array without conversion. Is this possible?

[rusty1s]

I think I fixed this in #7629 :)

[jagandecapri]

Hi @rusty1s ,

Are the fixes in the master branch? I tried installing the package using pip install git+https://github.com/pyg-team/pytorch_geometric.git and executed the code that I stated previously. The Data object still seems to be "converted" to a Numpy array.

[rusty1s]

Mh, looks like I checked in an incomplete fix. This does only resolve the issue if you remove the __len__ function from the Data object. I am not sure why adding __len__ leads numpy to think that it should convert this object to a dictionary. Need to investigate :(

[jagandecapri]

Noted. I found a workaround following https://stackoverflow.com/questions/43722023/can-i-prevent-numpy-array-from-casting-the-elements-as-numpy-arrays to allocate an empty Numpy array and then populating it with the Data object which doesn't convert the Data object to Numpy array. Sample code as below. Hope it helps someone. 🙂

import torch
import numpy as np
from torch_geometric.data import Data

x = torch.randn(10, 1)
edge_index = torch.randint(10, size=(2, 20))

a = Data(x=x, edge_index=edge_index)
b = np.empty((1), dtype=object)
b[:] = [a]
print(b)
print(f"Type: {type(b[0])}")

Output:

[Data(x=[10, 1], edge_index=[2, 20])]
Type: <class 'torch_geometric.data.data.Data'>

[rusty1s]

That's cool, thanks!