RWSENodeEncoder, KER_DIM_PE(lrgb文件中的encoders文件中的kernel.py)

该代码实现了一个基于核的节点编码器 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
评论
添加红包

请填写红包祝福语或标题

红包个数最小为10个

红包金额最低5元

当前余额3.43前往充值 >
需支付:10.00
成就一亿技术人!
领取后你会自动成为博主和红包主的粉丝 规则
hope_wisdom
发出的红包
实付
使用余额支付
点击重新获取
扫码支付
钱包余额 0

抵扣说明:

1.余额是钱包充值的虚拟货币,按照1:1的比例进行支付金额的抵扣。
2.余额无法直接购买下载,可以购买VIP、付费专栏及课程。

余额充值