该代码实现了一个基于核的节点编码器 KernelPENodeEncoder
,用于在图神经网络中将特定的核函数编码(例如随机游走结构编码 RWSE)与节点特征相结合。通过将预先计算的核统计信息(如 RWSE 等)与原始节点特征结合,该编码器可以帮助模型捕捉图中节点的结构信息。该代码还定义了一个派生类 RWSENodeEncoder
,这是一个具体实现,用于对节点执行随机游走结构编码。
编码器允许使用不同的模型(例如线性模型或 MLP)对核统计信息进行处理,并且支持根据需要对节点特征进行线性扩展。最终,节点特征和核特征被拼接成一个完整的节点嵌入。
from lrgb.encoders.kernel import RWSENodeEncoder, KER_DIM_PE
import torch
import torch.nn as nnKER_DIM_PE = 28
NUM_RW_STEPS = 20
MODEL = 'Linear'
LAYERS = 3
RAW_NORM_TYPE = 'BatchNorm'
PASS_AS_VAR = Falseclass KernelPENodeEncoder(torch.nn.Module):"""Configurable kernel-based Positional Encoding node encoder.The choice of which kernel-based statistics to use is configurable throughsetting of `kernel_type`. Based on this, the appropriate config is selected,and also the appropriate variable with precomputed kernel stats is thenselected 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 linearlyprojected to (dim_emb - dim_pe) size and the concatenated with PE.Args:dim_emb: Size of final node embeddingexpand_x: Expand node features `x` from dim_in to (dim_emb - dim_pe)"""kernel_type = None # Instantiated type of the KernelPE, e.g. RWSEdef __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 embeddingnum_rw_steps = NUM_RW_STEPSmodel_type = MODEL.lower() # Encoder NN model type for PEsn_layers = LAYERS # Num. layers in PE encoder modelnorm_type = RAW_NORM_TYPE.lower() # Raw PE normalization layer typeself.pass_as_var = PASS_AS_VAR # Pass PE also as a separate variableif dim_emb - dim_pe < 0: # formerly 1, but you could have zero feature sizeraise 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 > 0if norm_type == 'batchnorm':self.raw_norm = nn.BatchNorm1d(num_rw_steps)else:self.raw_norm = Noneactivation = 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