该代码实现了一个基于核的节点编码器 KernelPENodeEncoder
,用于在图神经网络中将特定的核函数编码(例如随机游走结构编码 RWSE)与节点特征相结合。通过将预先计算的核统计信息(如 RWSE 等)与原始节点特征结合,该编码器可以帮助模型捕捉图中节点的结构信息。该代码还定义了一个派生类 RWSENodeEncoder
,这是一个具体实现,用于对节点执行随机游走结构编码。
编码器允许使用不同的模型(例如线性模型或 MLP)对核统计信息进行处理,并且支持根据需要对节点特征进行线性扩展。最终,节点特征和核特征被拼接成一个完整的节点嵌入。
from lrgb.encoders.kernel import RWSENodeEncoder, KER_DIM_PE
import torch
import torch.nn as nn
KER_DIM_PE = 28
NUM_RW_STEPS = 20
MODEL = 'Linear'
LAYERS = 3
RAW_NORM_TYPE = 'BatchNorm'
PASS_AS_VAR = False
class KernelPENodeEncoder(torch.nn.Module):
"""Configurable kernel-based Positional Encoding node encoder.
The choice of which kernel-based statistics to use is configurable through
setting of `kernel_type`. Based on this, the appropriate config is selected,
and also the appropriate variable with precomputed kernel stats is then
selected from PyG Data graphs in `forward` function.
E.g., supported are 'RWSE', 'HKdiagSE', 'ElstaticSE'.
PE of size `dim_pe` will get appended to each node feature vector.
If `expand_x` set True, original node features will be first linearly
projected to (dim_emb - dim_pe) size and the concatenated with PE.
Args:
dim_emb: Size of final node embedding
expand_x: Expand node features `x` from dim_in to (dim_emb - dim_pe)
"""
kernel_type = None # Instantiated type of the KernelPE, e.g. RWSE
def __init__(self, dim_in, dim_emb, expand_x=True):
super().__init__()
if self.kernel_type is None:
raise ValueError(f"{self.__class__.__name__} has to be "
f"preconfigured by setting 'kernel_type' class"
f"variable before calling the constructor.")
dim_pe = KER_DIM_PE # Size of the kernel-based PE embedding
num_rw_steps = NUM_RW_STEPS
model_type = MODEL.lower() # Encoder NN model type for PEs
n_layers = LAYERS # Num. layers in PE encoder model
norm_type = RAW_NORM_TYPE.lower() # Raw PE normalization layer type
self.pass_as_var = PASS_AS_VAR # Pass PE also as a separate variable
if dim_emb - dim_pe < 0: # formerly 1, but you could have zero feature size
raise ValueError(f"PE dim size {dim_pe} is too large for "
f"desired embedding size of {dim_emb}.")
if expand_x and dim_emb - dim_pe > 0:
self.linear_x = nn.Linear(dim_in, dim_emb - dim_pe)
self.expand_x = expand_x and dim_emb - dim_pe > 0
if norm_type == 'batchnorm':
self.raw_norm = nn.BatchNorm1d(num_rw_steps)
else:
self.raw_norm = None
activation = nn.ReLU # register.act_dict[cfg.gnn.act]
if model_type == 'mlp':
layers = []
if n_layers == 1:
layers.append(nn.Linear(num_rw_steps, dim_pe))
layers.append(activation())
else:
layers.append(nn.Linear(num_rw_steps, 2 * dim_pe))
layers.append(activation())
for _ in range(n_layers - 2):
layers.append(nn.Linear(2 * dim_pe, 2 * dim_pe))
layers.append(activation())
layers.append(nn.Linear(2 * dim_pe, dim_pe))
layers.append(activation())
self.pe_encoder = nn.Sequential(*layers)
elif model_type == 'linear':
self.pe_encoder = nn.Linear(num_rw_steps, dim_pe)
else:
raise ValueError(f"{self.__class__.__name__}: Does not support "
f"'{model_type}' e