How to subclass Data class

[adrianomartinelli]

Hi,
I am trying to subclass the data class and adding some custom attributes and functionalities. However, I run into problems with the retrieval from the Globalstorage.

First of all, I would expect this snippet to set the attribute new_attr to None but it does not. Only if new_attr is not None

new_attr = None
d = Data(new_attr=new_attr)
d.new_attr

The overall goal is to add some validation steps to (new) attributes. In the end I would like to extend the Data class similar to

class MyData(Data):
   """Extends the torch_geometrics data representation"""

    def __init__(self, my_attr1=None, my_attr2=None, **kwargs):
        super().__init__(**kwargs)
        # self.my_attr1 = my_attr1
        # self.my_attr2 = my_attr2

        for k,v in zip(['my_attr1', 'my_attr2'], [my_attr1, my_attr2]):
            setattr(self, k, v)

    @property
    def my_attr1(self):
        return self._my_attr1

    @my_attr1.setter
    def my_attr1(self, value):
        if value is None:
            self._my_attr1 = None
        else:
            self._my_attr1 = torch.tensor(value)


d = MyData()
d.x
d.my_attr1

But neither approach (with setattr nor the standard way to set properties) works but results in a ttributeError: 'GlobalStorage' object has no attribute 'my_attr1' error.

What would you suggest to achieve this goal?

Thanks,
Adriano

[rusty1s]

We currently treat the setting of attributes to None as removing them from the data object. One way to resolve this is to do:

my_attr = data.get('my_attr', None)

[rusty1s]

What do you mean with validation steps? In general, it is recommended to always set attributes via __setattr__ in PyG - it just passes the values down to the underlying store. If you want to avoid that, you can assign it directly to __dict__ (not recommended):

data.__dict__["my_property"] = ...
print(data.my_property)

[adrianomartinelli]

With validation I mean for example that I want to ensure that an attribute holds only certain datatypes. In the example class below I want to enforce dtype=int.

class MyData(Data):

    def __init__(self, *args, bar=1.0, **kwargs):
        super().__init__(*args, **kwargs)
        self.bar = bar

    @property
    def bar(self):
        return self._store['bar']

    # this setter method is never accessed as __setattr__ takes precedence
    @bar.setter
    def bar(self, value):
        print('set bar value')
        # add some validation logic like enforcing type int
        self._store['bar'] = np.asarray(value, dtype=int)

I think I would need to overwrite __setattr__ now to handle such cases.

Might I ask what the idea behind the stores is?

Best,
Adriano

[rusty1s]

Interesting. Thanks for clarifying. The reason we override __setattr__ is to easily allow storing data in underlying stores (GlobalStorage in Data, and NodeStorage/EdgeStorage/GlobalStorage in HeteroData). For now, I think the only way to handle is to override __setattr__ by yourself:

class MyData(Data):

    def __init__(self, *args, bar=1.0, **kwargs):
        super().__init__(*args, **kwargs)
        self.bar = bar

    @property
    def bar(self):
        return self._store['bar']

    def __setattr__(self, key, value):
        if key == 'bar':
            self._store['bar'] = np.asarray(value, dtype=int)
        else:
            super().__setattr__(key, value)

Let me know if you can think of a more elegant solution.

[adrianomartinelli]

great thanks! What do you think about this solution?

• if a property setter method is defined we use it
• else default to Data.__setattr__ method

class Graph(Data):
    """Extends the torch_geometrics data representation"""

    def __init__(self, x: OptTensor = None, edge_index: OptTensor = None,
                 edge_attr: OptTensor = None, y: OptTensor = None,
                 pos: OptTensor = None, **kwargs):
        super().__init__(x, edge_index, edge_attr, y, pos, **kwargs)
        self.my_attr = 1  # stored in instance._store.__dict__['_my_attr']
        self.my_attr2 = 2  # stored in instance._store['my_attr2]

    @property
    def my_attr(self):
        return self._my_attr

    @my_attr.setter
    def my_attr(self, value):
        self._my_attr = torch.tensor(value, dtype=torch.float)

    def __setattr__(self, key, value):
        """overwrite to increase interoperability

        `__setattr__ takes precedence over properties and __dict__ <ref>`_. We ensure that, if the class defines
        a property, it is used, instead of executing pytorch_geometric default behaviour and saving the attribute in
        the `_store` of :cls:`~torch_geometric.data.Data`.

        .. _ref: https://stackoverflow.com/questions/15750522/class-properties-and-setattr
        """

        propobj = getattr(self.__class__, key, None)
        if propobj is None or propobj.fset is None:
            super().__setattr__(key, value)
        else:
            propobj.fset(self, value)

I did not extensively test this implementation yet but it seems to work so far.
Out of curiosity what is the benefit of using the stores? Is this to ensure separation of the data?

Cheers,
Adriano

[rusty1s]

I think this is a pretty cool solution. Happy to accept a PR on this one.

Separation of data is one good reason. In addition, overriding __setattr__ gives us a lot of flexibility on how we interact with HeteroData objects, e.g., by simply calling data.x_dict to get a dictionary of feature matrices across all types.

[rusty1s]

Fixed via #4682