Referencing the original node embeddings inside Aggregation in GCNConv

[MichaelHopwood]

mu and std aggregations have been set up in this package but I was hoping to get a z=(x-mu)/std where x is the embeddings of the current node.

How do we find x though? I see that an index is passed in to show which representations should be aggregated to which node. In other words, how can we reference the values of the seed nodes?

import argparse
import os.path as osp

import torch
import torch.nn.functional as F

import torch_geometric.transforms as T
from torch_geometric.datasets import Planetoid
from torch_geometric.logging import init_wandb, log
from torch_geometric.nn import GCNConv
from torch_geometric.nn.aggr.basic import StdAggregation, MeanAggregation, Aggregation, VarAggregation
import copy

from typing import Optional
from torch import Tensor
from torch.nn import Parameter

class Args:
    def __init__(self, args_dict):
        for key in args_dict:
            setattr(self, key, args_dict[key])

args_dict = {
    'dataset': 'Cora',
    'hidden_channels': 16,
    'lr': 0.01,
    'epochs': 200,
    'use_gdc': False,
    'wandb': False,
}
args = Args(args_dict)

device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
init_wandb(name=f'GCN-{args.dataset}', lr=args.lr, epochs=args.epochs,
           hidden_channels=args.hidden_channels, device=device)

path = osp.join('data-benchmark', 'Planetoid')
dataset = Planetoid(path, args.dataset, transform=T.NormalizeFeatures())
data = dataset[0]

if args.use_gdc:
    transform = T.GDC(
        self_loop_weight=1,
        normalization_in='sym',
        normalization_out='col',
        diffusion_kwargs=dict(method='ppr', alpha=0.05),
        sparsification_kwargs=dict(method='topk', k=128, dim=0),
        exact=True,
    )
    data = transform(data)

class ZScoreAggregation(Aggregation):
    def __init__(self):
        super().__init__()
        self.std_aggr = StdAggregation()
        self.mean_aggr = MeanAggregation()

    def forward(self, x: Tensor, index: Optional[Tensor] = None,
                ptr: Optional[Tensor] = None, dim_size: Optional[int] = None,
                dim: int = -2) -> Tensor:
        mu = self.mean_aggr(x, index, ptr, dim_size, dim)
        std = self.std_aggr(x, index, ptr, dim_size, dim)
       
        print(x.shape)
        print(mu.shape)
        print(std.shape)
        return (x - mu) / std

aggr = ZScoreAggregation()
   
class GCN(torch.nn.Module):
    def __init__(self, in_channels, hidden_channels, out_channels):
        super().__init__()
        self.conv1 = GCNConv(in_channels, hidden_channels, cached=True,
                             normalize=not args.use_gdc, aggr=copy.deepcopy(aggr))
        self.conv2 = GCNConv(hidden_channels, out_channels, cached=True,
                             normalize=not args.use_gdc, aggr=copy.deepcopy(aggr))

    def forward(self, x, edge_index, edge_weight=None):
        x = F.dropout(x, p=0.5, training=self.training)
        x = self.conv1(x, edge_index, edge_weight).relu()
        x = F.dropout(x, p=0.5, training=self.training)
        x = self.conv2(x, edge_index, edge_weight)
        return x


model = GCN(dataset.num_features, args.hidden_channels, dataset.num_classes)
model, data = model.to(device), data.to(device)
optimizer = torch.optim.Adam([
    dict(params=model.conv1.parameters(), weight_decay=5e-4),
    dict(params=model.conv2.parameters(), weight_decay=0)
], lr=args.lr)  # Only perform weight-decay on first convolution.


def train():
    model.train()
    optimizer.zero_grad()
    out = model(data.x, data.edge_index, data.edge_weight)
    loss = F.cross_entropy(out[data.train_mask], data.y[data.train_mask])
    loss.backward()
    optimizer.step()
    return float(loss)


@torch.no_grad()
def test():
    model.eval()
    pred = model(data.x, data.edge_index, data.edge_weight).argmax(dim=-1)

    accs = []
    for mask in [data.train_mask, data.val_mask, data.test_mask]:
        accs.append(int((pred[mask] == data.y[mask]).sum()) / int(mask.sum()))
    return accs


best_val_acc = final_test_acc = 0
for epoch in range(1, args.epochs + 1):
    loss = train()
    train_acc, val_acc, tmp_test_acc = test()
    if val_acc > best_val_acc:
        best_val_acc = val_acc
        test_acc = tmp_test_acc

Here is my requirements.txt:

numpy
pandas
tqdm
xgboost
scipy
scikit-learn
networkx
boto3
# automatically visualizing model performance
matplotlib
seaborn

--extra-index-url https://download.pytorch.org/whl/cpu
torch==1.12.1+cpu

--extra-index-url https://download.pytorch.org/whl/cpu
torchvision==0.13.1+cpu

--extra-index-url https://download.pytorch.org/whl/cpu
torchaudio==0.12.1+cpu

-f https://data.pyg.org/whl/torch-1.12.0+cpu.html
torch-scatter==2.0.9

-f https://data.pyg.org/whl/torch-1.12.0+cpu.html
torch-sparse==0.6.15

-f https://data.pyg.org/whl/torch-1.12.0+cpu.html
torch-cluster==1.6.0

-f https://data.pyg.org/whl/torch-1.12.0+cpu.html
torch-spline-conv==1.2.1

#-f https://data.pyg.org/whl/torch-1.12.0+cpu.html
#torch-geometric==2.1.0.post1

git+https://github.com/pyg-team/pytorch_geometric.git

[EdisonLeeeee]

It's not straightforward to do so in PyG, but you can achieve this with some necessary modifications:

from torch_geometric.nn import GCNConv
from torch_geometric.nn.aggr import MultiAggregation

class MyAggr(MultiAggregation):
    
    def combine(self, inputs):
        # inputs hold the outputs aggregated by `mean` and `std`
        return inputs

class MyConv(GCNConv):
    def update(self, inputs, x):
        mean, std = inputs
        return (x - mean) / std

aggr = MyAggr([MeanAggregation(), StdAggregation()])
MyConv(..., aggr=aggr)