quaterion.train.trainable_model 模块

class TrainableModel(*args: Any, **kwargs: Any)

基类：LightningModule, CacheMixin

用于训练模型的基类。

TrainableModel 用于描述模型的哪些组件以及如何进行训练。

它从构建块（如 Encoder, EncoderHead 等）组装模型。

┌─────────┐ ┌─────────┐ ┌─────────┐
│Encoder 1│ │Encoder 2│ │Encoder 3│
└────┬────┘ └────┬────┘ └────┬────┘
     │           │           │
     └────────┐  │  ┌────────┘
              │  │  │
          ┌───┴──┴──┴───┐
          │   concat    │
          └──────┬──────┘
                 │
          ┌──────┴──────┐
          │    Head     │
          └─────────────┘

TrainableModel 还处理大部分训练过程的例程：训练和验证步骤、张量设备管理、日志记录等等。大多数训练例程继承自 LightningModule，它是 TrainableModel 的直接父类。

要训练模型，您需要继承 TrainableModel 并实现所需的方法和属性。

最小示例

class ExampleModel(TrainableModel):
    def __init__(self, lr=10e-5, *args, **kwargs):
        self.lr = lr
        super().__init__(*args, **kwargs)

    # backbone of the model
    def configure_encoders(self):
        return YourAwesomeEncoder()

    # top layer of the model
    def configure_head(self, input_embedding_size: int):
        return SkipConnectionHead(input_embedding_size)

    def configure_optimizers(self):
        return Adam(self.model.parameters(), lr=self.lr)

    def configure_loss(self):
        return ContrastiveLoss()

configure_caches() → CacheConfig | None

提供缓存配置的方法

使用此方法定义哪些编码器应缓存计算出的 embeddings 以及它们应使用何种缓存。

返回值：

Optional[CacheConfig] – 如果提供，则应用缓存配置，否则为 None。

示例

不使用缓存（默认）:

return None

为所有不可训练的编码器自动配置缓存:

return CacheConfig(CacheType.AUTO)

为每个编码器单独指定缓存类型:

return CacheConfig(mapping={
        "text_encoder": CacheType.GPU,
        # Store cache in GPU for `text_encoder`
        "image_encoder": CacheType.CPU
        # Store cache in RAM for `image_encoder`
    }
)

指定缓存对象消歧的键:

return CacheConfig(
    cache_type=CacheType.AUTO,
    key_extractors={"text_encoder": hash}
)

如果您想提供更复杂的方式来存储缓存向量和原始对象之间的关联，此函数可能会很有用。如果未指定键，默认将使用数据集中项目的编号。

configure_encoders() → Encoder | Dict[str, Encoder]

提供编码器配置的方法

使用此函数定义编码器的初始状态。此函数应用于在训练前以及检查点加载期间为编码器分配初始值。

返回值：

Union[Encoder, Dict[str, Encoder]]：将被分配给 DEFAULT_ENCODER_KEY 的编码器实例，或名称与编码器的映射。

configure_head(input_embedding_size: int) → EncoderHead

使用此函数为模型的 head 层定义初始状态。

参数：

input_embedding_size – 编码器产生的 embeddings 大小

返回值：

EncoderHead – 将添加到模型顶部的 head

configure_loss() → SimilarityLoss

配置要使用的损失函数的方法。

configure_metrics() → AttachedMetric | List[AttachedMetric]

配置训练过程的批次级指标的方法

使用此方法将批次级指标附加到训练过程中。提供的指标必须类似于 PairMetric 或 GroupMetric

返回值：

Union[AttachedMetric, List[AttachedMetric]] - 附加到模型的指标

示例

return [
    AttachedMetric(
        "RetrievalPrecision",
        RetrievalPrecision(k=1),
        prog_bar=True,
        on_epoch=True,
    ),
    AttachedMetric(
        "RetrievalReciprocalRank",
        RetrievalReciprocalRank(),
        prog_bar=True,
    ),
]

configure_xbm() → XbmConfig

启用和配置跨批次内存 (XBM) 的方法。

XBM 是一种方法，依赖于训练过程中 embeddings “缓慢漂移”的思想。它将最近的 N 个 embeddings 和目标值保存在一个环形缓冲区中，其中 N 远大于批次大小。然后，它使用此缓冲区中的值计算缩放损失并将其添加到常规损失中。这使得能够挖掘大量困难负样本。

更多详情请参见论文：https://arxiv.org/pdf/1912.06798.pdf

要在训练过程中启用它，您必须返回 XbmConfig 的实例。默认返回值为 None，即不应用 XBM。

注意

XBM 目前仅支持 GroupLoss 实例。

process_results(embeddings: Tensor, targets: Dict[str, Any], batch_idx: int, stage: TrainStage, **kwargs)

提供 embeddings 附加评估的方法。

参数：

• embeddings – 形状: (batch_size, embedding_size) - 模型的输出。

• targets – 批次目标整理的输出。

• batch_idx – 正在处理的批次的 ID。

• stage – 训练、验证或测试阶段。

save_servable(path: str)

保存模型用于服务，独立于 Pytorch Lightning

参数：

path – 保存路径

setup_dataloader(dataloader: SimilarityDataLoader)

为编码器特定设置设置数据加载器，设置编码器特定整理函数

每个编码器都有自己独特的方式将记录列表转换为 NN 兼容格式。这些转换通常在数据预处理步骤中完成。

property loss: SimilarityLoss

获取要使用的损失函数的属性。

property model: SimilarityModel

要训练的原始模型

返回值：

SimilarityModel – 要训练的模型