开发者指南

备注

为了更好地理解以下指南，您可以先查看我们的论文来了解整体思路。

GLUE是模块化的框架，并且能以多种方式进行扩展。

下面我们将介绍该框架的主要组件，以及如何扩展现有框架。

主要组件

GLUE模型由四个主要组件组成（全为PyTorch 模块）:

每个域有一个数据编码器 —— 输入数据 \(x\) ，返回该数据的后验分布（细胞嵌入） \(q(u|x)\)
每个域有一个数据解码器 —— 输入细胞嵌入 \(u\) 和特征嵌入 \(v\) ，返回该数据的似然分布
一个图编码器 —— 输入图 \(\mathcal{G}\) （包括边索引、边权重和边符号），返回该图的后验分布（特征嵌入） \(q(v|\mathcal{G})\)
一个图解码器 —— 输入特征嵌入 \(v\) 和包括边索引和边符号的边子集，返回该边子集的似然分布，该分布被用作图的似然估计 \(p(\mathcal{G}|v)\)

目前，这些组件的实现都位于 scglue.models.sc 中，也可以在这个模块中添加新的扩展。

实际的模块接口可能与上面总结略有不同，如库大小的归一化和批量效应会有额外的考虑。详情可见下文。

支持新的数据模式

一个简单的扩展是添加新的数据编码器和解码器以支持更多的数据模式。

定义编码器

数据编码器继承自 DataEncoder 类。编码器的主要部分是MLP（多层感知器），已经在 DataEncoder 中实现，并且保留了两个可自定义的抽象方法：

compute_l 方法能根据输入数据计算出库的大小。
在输入MLP之前， normalize 方法能对输入数据(可能连同计算得到的库大小)进行规范化处理。

下面是一个负二项式数据编码器的例子。输入原始计数，通过计算每个单元格的计数之和来计算库大小，然后行规范化到10000，最后对规范化后的数据进行对数转换。

class NBDataEncoder(DataEncoder):

  TOTAL_COUNT = 1e4

  def compute_l(self, x: torch.Tensor) -> torch.Tensor:
      return x.sum(dim=1, keepdim=True)

  def normalize(
          self, x: torch.Tensor, l: torch.Tensor
  ) -> torch.Tensor:
      return (x * (self.TOTAL_COUNT / l)).log1p()

为适应新的数据模式，您可以通过实现这两个方法来定义新的编码器类。

定义解码器

数据解码器继承自 DataDecoder 类，它定义了构造器的接口和抽象方法 forward。

构造器接受两个输入：输出维度 out_features 和批次数量 n_batches (这里的批次是相对于批次效应而言)。 forward 方法接受四个输入：

u 表示细胞嵌入
v 表示特征嵌入
b 表示批次索引
l 表示编码器计算得到的库大小

返回该数据的似然分布。

下面是一个负二项式数据解码器的例子。它包括三个可训练的参数 scale_lin, bias 和 log_theta (你也可以根据需要定义自己的参数)：

scale_lin 给出了soft-plus转换后的比例参数 \(\alpha\) ，保证正向性
bias 表示偏置参数 \(\beta\)
log_theta 表示负二项式的反离散参数 :math:` heta` 的对数

这三个参数都被定义为批次特定的参数(每个批次的参数由不同的行定义)。

通过缩放、移动细胞嵌入和特征嵌入的内积得到logit_mu，logit_mu经过softmax函数并与库大小相乘得到负二项式的平均数 mu，返回一个负二项分布。

class NBDataDecoder(DataDecoder):

    def __init__(self, out_features: int, n_batches: int = 1) -> None:
        super().__init__(out_features, n_batches=n_batches)
        self.scale_lin = torch.nn.Parameter(torch.zeros(n_batches, out_features))
        self.bias = torch.nn.Parameter(torch.zeros(n_batches, out_features))
        self.log_theta = torch.nn.Parameter(torch.zeros(n_batches, out_features))

    def forward(
            self, u: torch.Tensor, v: torch.Tensor,
            b: torch.Tensor, l: torch.Tensor
    ) -> D.NegativeBinomial:
        scale = F.softplus(self.scale_lin[b])
        logit_mu = scale * (u @ v.t()) + self.bias[b]
        mu = F.softmax(logit_mu, dim=1) * l
        log_theta = self.log_theta[b]
        return D.NegativeBinomial(
            log_theta.exp(),
            logits=(mu + EPS).log() - log_theta
        )

注意批次索引 b 作为行索引进入 scale_lin, bias 和 log_theta。

为生成适应新数据模式的似然分布，您可以通过实现 forward 方法来定义新的解码器类。

非标准分布可以在 scglue.models.prob 中定义。

注册自定义的编码器和解码器

最后，通过 scglue.models.scglue.register_prob_model 函数在一个新的 “prob_model” 下注册自定义的编码器和解码器，并在 configure_dataset 中用匹配的 prob_model 激活。

其他类型的扩展？

如果您对以其他方式扩展模型感兴趣，请在 Github 上开启一个问题。

欢迎投稿！

如果您希望您的扩展被包含在框架中，请务必在 Github 上提交 “pull” 请求。