Get node embeddings from trained GraphGym model

[nilskre]

Hello,

thanks for this amazing project!

After using GraphGym to train models I would like to get the node embeddings and further analyze them.
My approach would be to load the trained model from GraphGym (stored e.g. under graphgym/results/example_node/0/ckpt/epoch=199-step=200.ckpt) and use this model to get the node embeddings.

Normally loading a Pytorch model would be possible by:
model = TheModelClass(*args, **kwargs)
model.load_state_dict(torch.load(PATH))

How is it possible to load the saved GraphGym model (how is TheModelClass() initialized?) and use it to generate the node embeddings?

Thanks in advance!

[rusty1s]

It would be

model = create_model()
model.load_state_dict(torch.load(PATH))

(see https://github.com/pyg-team/pytorch_geometric/blob/master/graphgym/main.py).

[nilskre]

Alright, thanks! I was now able to load the model again with:

import torch
from torch_geometric.graphgym.model_builder import create_model
from torch_geometric.graphgym.config import cfg, load_cfg
from argparse import Namespace

args = Namespace(cfg_file='pytorch_geometric/graphgym/configs/pyg/example_node.yaml', repeat=1, mark_done=False, opts=[])
load_cfg(cfg, args)
model = create_model()
model.load_state_dict(torch.load(path_to_model)['state_dict'])

Resulting in (the standard model from the GraphGym single example):

GraphGymModule(
  (model): GNN(
    (encoder): FeatureEncoder()
    (mp): GNNStackStage(
      (layer0): GeneralLayer(
        (layer): GCNConv(
          (model): GCNConv(1, 16)
        )
        (post_layer): Sequential(
          (0): Dropout(p=0.1, inplace=False)
          (1): PReLU(num_parameters=1)
        )
      )
      (layer1): GeneralLayer(
        (layer): GCNConv(
          (model): GCNConv(16, 16)
        )
        (post_layer): Sequential(
          (0): Dropout(p=0.1, inplace=False)
          (1): PReLU(num_parameters=1)
        )
      )
    )
    (post_mp): GNNNodeHead(
      (layer_post_mp): MLP(
        (model): Sequential(
          (0): Linear(
            (model): Linear(16, 1, bias=True)
          )
        )
      )
    )
  )
)

Sorry to bother you again, but one follow-up question has now arisen:
I would like to use the model to get the node embeddings. Therefore I have passed my dataset (Data(x=[3013686, 1], edge_index=[2, 6364814], y=[3013686], train_mask=[3013686], val_mask=[3013686], test_mask=[3013686])) to my model by

out = model(data)

This leads to AttributeError: 'GlobalStorage' object has no attribute 'split'

When trying to call it similar to the tutorial with
out = model(data.x, data.edge_index), I get TypeError: forward() takes 2 positional arguments but 3 were given

Where is my fault?

Thanks a lot in advance!

[rusty1s]

Thanks for reporting. I fixed this in https://github.com/pyg-team/pytorch_geometric/pull/7327/files. Sorry for any inconvenience.

[nilskre]

No worries, thanks a lot for the fast fix! Now calling the trained model works without any problem.

But when looking at the model output one hopefully final question has arisen:

At the moment it seems to be the case, that the model returns the predicted labels and the ground truth labels. How is it possible to also access the learned node embeddings? (Final goal would be to learn node embeddings fully unsupervised and then compare only the final node embeddings to the ground truth labels).

One option, I could imagine at the moment is to implement a custom node head, which also returns the embeddings. But is there maybe a more simple solution, I overlook at the moment?
I also found the config option view_emb, but I am not sure, if this has any effect.

Thanks again!

For debugging purposes the config.yaml:

accelerator: cpu
benchmark: false
bn:
  eps: 1.0e-05
  mom: 0.1
cfg_dest: config.yaml
custom_metrics: []
dataset:
  cache_load: false
  cache_save: false
  dir: ./datasets
  edge_dim: 128
  edge_encoder: false
  edge_encoder_bn: true
  edge_encoder_name: Bond
  edge_message_ratio: 0.8
  edge_negative_sampling_ratio: 1.0
  edge_train_mode: all
  encoder: true
  encoder_bn: true
  encoder_dim: 128
  encoder_name: db
  format: PyG
  label_column: none
  label_table: none
  location: local
  name: ckg_subset
  node_encoder: false
  node_encoder_bn: true
  node_encoder_name: Atom
  remove_feature: false
  resample_disjoint: false
  resample_negative: false
  shuffle_split: true
  split:
  - 0.8
  - 0.1
  - 0.1
  split_mode: random
  task: node
  task_type: regression
  to_undirected: false
  transductive: true
  transform: none
  tu_simple: true
devices: 1
gnn:
  act: prelu
  agg: mean
  att_final_linear: false
  att_final_linear_bn: false
  att_heads: 1
  batchnorm: false
  clear_feature: true
  dim_inner: 16
  dropout: 0.1
  head: node
  keep_edge: 0.5
  l2norm: true
  layer_type: gcnconv
  layers_mp: 2
  layers_post_mp: 1
  layers_pre_mp: 0
  msg_direction: single
  normalize_adj: false
  self_msg: concat
  skip_every: 1
  stage_type: stack
gpu_mem: false
mem:
  inplace: false
metric_agg: argmax
metric_best: auto
model:
  edge_decoding: dot
  graph_pooling: add
  loss_fun: mse
  match_upper: true
  size_average: mean
  thresh: 0.5
  type: gnn
num_threads: 6
num_workers: 0
optim:
  base_lr: 0.01
  lr_decay: 0.1
  max_epoch: 2
  momentum: 0.9
  optimizer: adam
  scheduler: cos
  steps:
  - 30
  - 60
  - 90
  weight_decay: 0.0005
out_dir: results/example_node
print: both
round: 4
run_dir: results
seed: 0
share:
  dim_in: 1
  dim_out: 1
  num_splits: 1
tensorboard_agg: true
tensorboard_each_run: false
train:
  auto_resume: false
  batch_size: 128
  ckpt_clean: true
  ckpt_period: 100
  enable_ckpt: true
  epoch_resume: -1
  eval_period: 1
  iter_per_epoch: 32
  neighbor_sizes:
  - 20
  - 15
  - 10
  - 5
  node_per_graph: 32
  radius: extend
  sample_node: false
  sampler: full_batch
  skip_train_eval: false
  walk_length: 4
val:
  node_per_graph: 32
  radius: extend
  sample_node: false
  sampler: full_batch
view_emb: true

[rusty1s]

Good question. I think the most elegant solution to this is to register a forward_hook from PyTorch on top of your model, which lets you get the embeddings prior to the model head. I think this is the best solution since one usually does not want to modify the model output for these cases.

[nilskre]

Thanks a lot! This is really an elegant solution.

For everybody with the same use case the code from here:

from typing import Dict, Iterable, Callable

class EmbeddingsExtractor(nn.Module):
    def __init__(self, model: nn.Module, layers: Iterable[str]):
        super().__init__()
        self.model = model
        self.layers = layers
        self._features = {layer: torch.empty(0) for layer in layers}

        for layer_id in layers:
            layer = dict([*self.model.named_modules()])[layer_id]
            layer.register_forward_hook(self.save_outputs_hook(layer_id))

    def save_outputs_hook(self, layer_id: str) -> Callable:
        def fn(_, __, output):
            self._features[layer_id] = output
        return fn

    def forward(self, x: torch.Tensor) -> Dict[str, torch.Tensor]:
        _ = self.model(x)
        return self._features

Select the layer from which you would like to get the output. To see all the layers: dict([*model.named_modules()]).keys()

Then get the node embeddings from this layer:

get_node_embeddings = EmbeddingsExtractor(model, layers=["model.mp.layer1.post_layer.1"])
embeddings_for_layers = get_node_embeddings(data)