跳到内容

添加组合器

组合器负责将一个或多个输入特征的输出组合成一个单一的组合表示,通常是一个向量,但也可能是向量序列或其他更高维度的张量。一个或多个输出特征将使用这种组合表示来生成预测。

用户可以在配置的 combiner 部分指定要使用的组合器;如果未指定组合器,将使用 concat 组合器。

回顾 ECD (编码器、组合器、解码器) 数据流架构:所有输入特征的输出流向组合器,组合器的输出流向所有输出特征。

+-----------+                      +-----------+
|Input      |                      | Output    |
|Feature 1  +-+                  +-+ Feature 1 + ---> Prediction 1
+-----------+ |                  | +-----------+
+-----------+ |   +----------+   | +-----------+
|...        +---> | Combiner +---> |...        +
+-----------+ |   +----------+   | +-----------+
+-----------+ |                  | +-----------+
|Input      +-+                  +-+ Output    |
|Feature N  |                      | Feature N + ---> Prediction N
+-----------+                      +-----------+

还需要注意一个额外的复杂性:输入特征可能输出向量,也可能输出向量序列。因此,组合器可能需要处理输出维度不同的混合输入特征。例如,SequenceConcatCombiner 通过要求所有输入序列具有相同的长度来解决这个问题。如果长度不一致,它会引发 ValueError 异常。SequenceConcatCombiner 在连接之前将非序列输入填充到序列长度,从而将所有输入特征作为相同长度的序列进行处理。

新的组合器应该在其文档字符串中明确说明是否支持序列输入,声明对序列长度、类型或维度的任何要求,并验证其输入特征。

本指南将概述如何通过添加新的组合器来扩展 Ludwig,并以 transformer 组合器作为模板。要添加新的组合器:

  1. 定义一个数据类来表示组合器模式。
  2. 创建一个继承自 ludwig.combiners.Combiner 或其子类的新组合器类。
  3. __init__ 方法中分配所有层和状态。
  4. def forward(self, inputs: Dict): 中实现组合器的前向传播。
  5. 添加测试。
  6. 将新的组合器添加到组合器注册表。

1. 定义组合器的模式

组合器模式是一个 dataclass(由 marshmallow_dataclass 模块重载),它必须扩展 BaseCombinerConfig。其属性是组合器的配置参数。所有字段都应具有类型和默认值。ludwig.schema.utils.py 模块提供了方便的方法来指定组合器配置的有效类型和范围。例如,TransformerCombiner 具有以下模式:

from typing import Optional, List, Dict, Any

# Main imports:
from marshmallow_dataclass import dataclass
from ludwig.schema import utils as schema_utils
from ludwig.schema.combiners.base import BaseCombinerConfig

@dataclass
class TransformerCombinerConfig(BaseCombinerConfig):
    num_layers: int = schema.PositiveInteger(default=1)
    hidden_size: int = schema.NonNegativeInteger(default=256)
    num_heads: int = schema.NonNegativeInteger(default=8)
    transformer_output_size: int = schema.NonNegativeInteger(default=256)
    dropout: float = schema.FloatRange(default=0.1, min=0, max=1)
    fc_layers: Optional[List[Dict[str, Any]]] = schema.DictList()
    num_fc_layers: int = schema.NonNegativeInteger(default=0)
    output_size: int = schema.PositiveInteger(default=256)
    use_bias: bool = True
    weights_initializer: Union[str, Dict] = \
        schema.InitializerOrDict(default="xavier_uniform")
    bias_initializer: Union[str, Dict] = \
        schema.InitializerOrDict(default="zeros")
    norm: Optional[str] = schema.StringOptions(["batch", "layer"])
    norm_params: Optional[dict] = schema.Dict()
    fc_activation: str = "relu"
    fc_dropout: float = schema.FloatRange(default=0.0, min=0, max=1)
    fc_residual: bool = False
    reduce_output: Optional[str] = schema.ReductionOptions(default="mean")

这个模式应该放在 ludwig/schema/combiners/ 目录下的独立文件中。为了方便在 Ludwig 的其他地方导入,您还可以将其作为导入添加到 ludwig/schema/combiners/__init__.py 文件中,如下所示:

from ludwig.schema.combiners.transformer import TransformerCombinerConfig  # noqa: F401

2. 添加新的组合器类

组合器的源代码位于 ludwig/combiners/。添加一个新的 Python 模块来声明一个新的组合器类。在本示例中,我们将展示如何实现一个简化版的 transformer 组合器,它将被定义在 transformer_combiner.py 文件中。

注意

目前,所有组合器都定义在 ludwig/combiners/combiners.py 文件中。然而,对于新的组合器,我们建议创建一个与新组合器类名称对应的新的 Python 模块。

@register_combiner(name="transformer")
class TransformerCombiner(Combiner):
    def __init__(
        self,
        input_features: Dict[str, InputFeature] = None,
        config: TransformerCombinerConfig = None,
        **kwargs
    ):
        super().__init__(input_features)
        self.name = "TransformerCombiner"

    def forward(
        self,
        inputs: Dict,
    ) -> Dict[str: torch.Tensor]:

    @staticmethod
    def get_schema_cls():
        return TransformerCombinerConfig

实现 @staticmethod def get_schema_cls(): 方法并返回您的配置模式的类名。

3. 实现构造函数

组合器构造函数将使用输入特征的字典和组合器配置进行初始化。构造函数必须将输入特征传递给父类构造函数,设置其 name 属性,然后创建自己的层和状态。

input_features 字典传递给构造函数,以便获取关于输入数量、大小和类型的信息。这可以决定组合器需要分配哪些资源。例如,transformer 组合器将其输入特征视为一个序列,其中序列长度就是特征数量。我们可以在这里确定序列长度为 self.sequence_size = len(self.input_features)

    def __init__(
        self,
        input_features: Dict[str, InputFeature] = None,
        config: TransformerCombinerConfig = None,
        **kwargs
    ):
        super().__init__(input_features)
        self.name = "TransformerCombiner"
        # ...
        self.sequence_size = len(self.input_features)

        self.transformer_stack = TransformerStack(
            input_size=config.hidden_size,
            sequence_size=self.sequence_size,
            hidden_size=config.hidden_size,
            num_heads=config.num_heads,
            output_size=config.transformer_output_size,
            num_layers=config.num_layers,
            dropout=config.dropout,
        )
        # ...

4. 实现 forward 方法

组合器的 forward 方法应将输入特征表示组合成一个单一的输出张量,该张量将被传递给输出特征解码器。输入中的每个键都是一个输入特征名称,相应的值是该输入特征输出的字典。每个特征输出字典保证包含一个 encoder_output 键,并可能根据编码器包含其他输出。

forward 返回一个将字符串映射到张量的字典,该字典必须包含一个 combiner_output 键。它还可以选择性地返回其他可能对输出特征解码、损失计算或解释有用的值。例如,TabNetCombiner 返回其稀疏注意力掩码(attention_masksaggregated_attention_masks),这有助于查看在每个预测步骤中哪些输入特征受到了关注。

例如,以下是 TransformerCombinerforward 方法的简化版本:

    def forward(
        self, inputs: Dict[str, Dict[str, torch.Tensor]]
    ) -> Dict[str, torch.Tensor]:
        encoder_outputs = [inputs[k]["encoder_output"] for k in inputs]

        # ================ Flatten ================
        batch_size = encoder_outputs[0].shape[0]
        encoder_outputs = [
            torch.reshape(eo, [batch_size, -1]) for eo in encoder_outputs
        ]

        # ================ Project & Concat ================
        projected = [
            self.projectors[i](eo) for i, eo in enumerate(encoder_outputs)
        ]
        hidden = torch.stack(projected)
        hidden = torch.permute(hidden, (1, 0, 2))

        # ================ Transformer Layers ================
        hidden = self.transformer_stack(hidden)

        # ================ Sequence Reduction ================
        if self.reduce_output is not None:
            hidden = self.reduce_sequence(hidden)
            hidden = self.fc_stack(hidden)

        return_data = {"combiner_output": hidden}
        return return_data

输入

  • inputs (Dict[str, Dict[str, torch.Tensor]]):一个字典,键为输入特征名称,值为该输入特征的输出字典。每个输入特征输出字典保证包含 encoder_output 键,并可能根据输入特征的编码器包含其他键值对。

返回值

  • (Dict[str, torch.Tensor]):一个字典,包含必需的 combiner_output 键,其值是组合器输出张量,以及任何其他可选的输出键值对。

5. 将新类添加到注册表

模型配置中的组合器名称与组合器类之间的映射是通过在组合器注册表中注册类来完成的。组合器注册表定义在 ludwig/schema/combiners/utils.py 文件中。要注册您的类,请在其类定义上一行添加 @register_combiner 装饰器,并指定组合器的名称:

@register_combiner(name="transformer")
class TransformerCombiner(Combiner):

6. 添加测试

tests/ludwig/combiners 目录中添加一个相应的单元测试模块,使用您的组合器模块名称并在前面加上 test_,例如 test_transformer_combiner.py

单元测试至少应确保:

  1. 组合器的前向传播对于其支持的所有特征类型都能成功。
  2. 当给定不支持的输入时,组合器按预期方式失败。(如果组合器支持所有输入特征类型,则跳过此项。)
  3. 在各种配置下,组合器产生正确类型和维度的输出。

使用 @pytest.mark.parametrize 通过不同的配置来参数化您的测试,同时测试边缘情况。

@pytest.mark.parametrize("output_size", [8, 16])
@pytest.mark.parametrize("transformer_output_size", [4, 12])
def test_transformer_combiner(
        encoder_outputs: tuple,
        transformer_output_size: int,
        output_size: int) -> None:
    encoder_outputs_dict, input_feature_dict = encoder_outputs

有关组合器测试的示例,请参阅 tests/ludwig/combiners/test_combiners.py

有关 Ludwig 中单元测试的更多详细信息,另请参阅 单元测试设计指南