Data Transform hack to apply Sklearn Feature Scaling on your custom Dataset class

[snknitin]

Hey PyG Team,

I've been working on Heterogenous Graphs for a while and I have multiple edges with edge attributes in my HeteroData(). Since there is no data transform that does feature scaling, and Standardization and pre-processing is currently expected to be handled by the user. Sometimes you might want your input embeddings or feature values to remain unchanged from raw files, and use transforms instead to Standardize or Normalize them based on use-case. Hardcoding the scaling in the Dataset process() method might pose a couple of issues, so I've built a custom transform that handles this.

You can refer to this discussion from last month to get a better idea. I finally got some time to refine and post my approach for this.

Quick Context

• Stats like mean, min, max, std have to be pre-computed in order to do ensure consistency across examples/graphs.
• If you have multiple graphs in your dataset, you need to compute stats from observing all edges and then re-scale individual graph edge features based on that.
• Multiple edges will have their own separate scaler that acts on those dimensions
• Transforms are applied on access and stats of the data need to be available for scaling, which means, you have to run your data twice, once without transforms, then compute stats, delete that data object and pass stats to the custom transform to redo the data but this time with feature scaling

Solution

In your process method when creating your own Dataset, just like you create

data['node1', 'to', 'node2'].edge_index = ...
data['node1', 'to', 'node2'].edge_attr = ...
data['node1', 'to', 'node2'].edge_label = ...

Add a

data['node1', 'to', 'node2'].edge_scaler = scaler

This way your edge_store will now have a scaler item, which won't be memory heavy. You can calculate this scaler within the Dataset class with a method like this -

def get_scalers(self,edge_cols):
    """
    Scaling the edge attributes based on the whole data for consistency
    we need to return scaler and perform the transform
    """

    def f(i):
        return pd.read_csv(i, usecols=edge_cols, low_memory=True)

    # Read edge_cols from all raw files
    df = pd.concat(map(f, self.raw_file_names))
    scaler = MinMaxScaler()
    scaler.fit(df)
    del df
    return scaler

Here I'm reading all my data for graphs from .csv files. I choose specific colums of the files that are going to be that specifc edge's features and fit the scaler from all the raw files. You can choose any scaling approach from sklearn.preprocessing. Now you can collate and process the graph.

Transform

@functional_transform('scale_edges')
class ScaleEdges(BaseTransform):
    r"""Column-normalizes the attributes given in :obj:`attrs` to sum-up to standardize or scale them.

    Mainly useful for edges, if the scalar for the collective data edges is stored
    as data[(node1,to,node2)].edge_scaler = scaler

    where the scaler can be StandardScaler or MinMaxScaler from scikit-learn that
    has been fit to the total dataset

    Args:
        attrs (List[str]): The names of attributes to normalize.
            (default: :obj:`["edge_attr"]`)
    """
    def __init__(self, attrs: List[str] = ["edge_attr"]):
        self.attrs = attrs

    def __call__(self, data: Union[Data, HeteroData]):
        for edge,store in list(zip(data.edge_types,data.edge_stores)):
            for key, value in store.items(*self.attrs):
                if "edge_scaler" in store.keys():
                    scaler = store["edge_scaler"]
                    # needs to be a float tensor since this is the edge_attr
                    value = torch.tensor(scaler.transform(value),dtype=torch.float)
                    store[key] = value
                    del store["edge_scaler"]
        return data

    def __repr__(self) -> str:
        return f'{self.__class__.__name__}()'

Now when you want to load your data and apply the transform, you can do it this way. Note that you would ideally want to run scaling before the other transforms like NormalizeFeatures. This can be modified to work on Node features as well.

transform = T.Compose([ ScaleEdges(attrs=["edge_attr"]),
                        T.ToUndirected(),
                        T.AddSelfLoops(),
                        T.NormalizeFeatures(attrs=["x", "edge_attr"]),
                        RevDelete()])

data = MyDataset(root,transform=transform)
data[0]

Bonus Transform

If you're curious about the RevDelete(), that is another transform I made to delete the edge_label that forms on the reverse edges in HeteroData objects after the ToUndirected() transform. It's a simple code to avoid manual effort

@functional_transform('revlabel_delete')
class RevDelete(BaseTransform):
    r"""Deletes the edge_label on reverse edges.
    Args:
        attrs (List[str]): The names of attributes to consider.
            (default: :obj:`["edge_attr"]`)
    """
    def  __init__(self, attrs: List[str] = ["edge_attr"]):
        self.attrs = attrs

    def __call__(self, data: Union[Data, HeteroData]):
        for edge,store in list(zip(data.edge_types,data.edge_stores)):
            if edge[1].startswith("rev"):
                if "edge_label" in store.keys():
                    del store["edge_label"]
        return data

    def __repr__(self) -> str:
        return f'{self.__class__.__name__}()'

Acknowledgement

Hope this is helpful to someone. I would love to know what you guys think and if there is a better way to do this. Let me know if I can contribute to the package in some way. I'd be stoked, if these could be refined and adopted into the built-in methods !

[rusty1s]

I think this is cool, thanks for sharing. We could think about integrating this as part of a SKLearnPreprocessTransform - a transform that gets initialized by taking a dataset in, computes the stats, and then applies them on every access. WDYT?

[snknitin]

Great ! I would love to. This would be my first time doing a PR. I'd be grateful for any guidance
Do you want me to fork, commit this, as is, from above, and start a PR for now, and to later build on comments ? I can work on this post-data_creation-processing sklearn transform idea over the weekend. Or, Do you want me to create one after that ?

I'm assuming you mean it would be something like this

data = T.ToUndirected(data)
data = T.SKLearnPreprocessTransform(data)

[rusty1s]

Yes, you would need to fork, push some initial draft and send a PR to PyG. We can improve it together. I think the high-level interface would be something like:

transform = SKLearnPreprocessTransform(dataset, some_other_options)
data = transform(data)

[snknitin]

Done !

[YoavLotem]

Hey @rusty1s !
This is an awesome capability.. did it mature?

[snknitin]

I worked on this a little bit in this PR with the hack I used for my project. It has the edge scaling transform but to make it a generic sklearn transform pipeline, it will need some modifications. I got busy last 3 months, so I didn't get a chance to continue. But I can pick it up again with some assistance