监督分类是机器学习中最广泛使用的训练目标之一，但并非所有任务都能如此定义。例如，
- 您的类别可能会快速变化——例如，随着时间推移可能会添加新类别，
- 您可能没有来自所有可能类别的样本，
- 在训练期间可能无法枚举所有可能的类别，
- 您可能有一个本质上不同的任务，例如搜索或检索。
所有这些问题都可以通过相似度学习高效解决。
注意：如果您不熟悉相似度学习概念，请查看 awesome-metric-learning 仓库以获取丰富的资源和用例示例。
然而，相似度学习也有其自身的难点，例如：
- 通常需要更大的批量大小，
- 更复杂的损失函数，
- 训练和推理之间架构的变化。
Quaterion 是一个为解决相似度学习中的此类问题而构建的微调框架。它使用 PyTorch Lightning 作为后端，其宣传口号是“花更多时间在研究上，更少时间在工程上”。这同样适用于 Quaterion，并且它包括：
- 可训练和服务化的模型类，
- 带注解的内置损失函数，以及当您需要更多功能时对 pytorch-metric-learning 的包装，
- 样本、数据集和数据加载器类，以便更轻松地处理相似度学习数据，
- 一种缓存机制，可实现更快的迭代和更少的内存占用。
深入了解 Quaterion
让我们分解一些重要模块
TrainableModel：它是pl.LightNingModule的子类，具有额外的钩子方法，例如configure_encoders、configure_head、configure_metrics等，用于定义训练和评估所需的各种对象——参见下方了解更多详情。SimilarityModel：一种仅用于推理的导出方法，旨在提高代码移植性并减少推理时的依赖。实际上，Quaterion 由两个包组成：quaterion_models：推理所需的包。quaterion：定义训练所需对象的包，也依赖于quaterion_models。
Encoder和EncoderHead：构成SimilarityModel的两个对象。在大多数情况下，您可以使用冻结的预训练编码器，例如来自torchvision的 ResNets，或来自transformers的语言模型，并在其顶部堆叠一个可训练的EncoderHead。quaterion_models提供了 多种EncoderHead实现，具有统一的 API，例如可配置的 dropout 值。您可以使用其中一种，也可以通过继承父类或在SequentialHead中轻松列出 PyTorch 模块来创建自己的EncoderHead。
Quaterion 还有其他对象，如距离函数、评估指标、评估器、便捷的数据集和数据加载器类，但这些大多数都是不言自明的。因此，本文将不再详细解释它们，以保持简洁。但是，您可以随时查看 文档 来了解更多信息。
本教程的重点是使用 Quaterion 逐步解决一个相似度学习问题。这也将帮助我们更好地理解上述对象在实际项目中如何协同工作。让我们开始浏览代码中的一些重要部分。
如果您想查看完整的源代码，可以在 Quaterion 仓库的 examples 目录下找到。
数据集
在本教程中，我们将使用 Stanford Cars 数据集。
Stanford Cars 数据集
它包含来自 196 个类别的 16185 张汽车图片，并被分割成训练集和测试集，几乎是五五开。然而，为了使其更有趣，我们将首先合并训练集和测试集，然后再次将其分割成两部分，使得 196 个类别中的一半进入训练集，另一半进入测试集。这将允许我们使用在训练阶段从未见过的新类别样本来测试我们的模型，这是监督分类无法实现但相似度学习可以做到的。
在以下代码中，该代码借用了 data.py：
get_datasets()函数执行上述分割任务。get_dataloaders()函数从训练集和测试集创建GroupSimilarityDataLoader实例。- 数据集是常规的 PyTorch 数据集，生成
SimilarityGroupSample实例。
注意：目前，Quaterion 有两种数据类型来表示数据集中的样本。要了解更多关于 SimilarityPairSample 的信息，请查看 NLP 教程。
import numpy as np
import os
import tqdm
from torch.utils.data import Dataset, Subset
from torchvision import datasets, transforms
from typing import Callable
from pytorch_lightning import seed_everything
from quaterion.dataset import (
GroupSimilarityDataLoader,
SimilarityGroupSample,
)
# set seed to deterministically sample train and test categories later on
seed_everything(seed=42)
# dataset will be downloaded to this directory under local directory
dataset_path = os.path.join(".", "torchvision", "datasets")
def get_datasets(input_size: int):
# Use Mean and std values for the ImageNet dataset as the base model was pretrained on it.
# taken from https://www.geeksforgeeks.org/how-to-normalize-images-in-pytorch/
mean = [0.485, 0.456, 0.406]
std = [0.229, 0.224, 0.225]
# create train and test transforms
transform = transforms.Compose(
[
transforms.Resize((input_size, input_size)),
transforms.ToTensor(),
transforms.Normalize(mean, std),
]
)
# we need to merge train and test splits into a full dataset first,
# and then we will split it to two subsets again with each one composed of distinct labels.
full_dataset = datasets.StanfordCars(
root=dataset_path, split="train", download=True
) + datasets.StanfordCars(root=dataset_path, split="test", download=True)
# full_dataset contains examples from 196 categories labeled with an integer from 0 to 195
# randomly sample half of it to be used for training
train_categories = np.random.choice(a=196, size=196 // 2, replace=False)
# get a list of labels for all samples in the dataset
labels_list = np.array([label for _, label in tqdm.tqdm(full_dataset)])
# get a mask for indices where label is included in train_categories
labels_mask = np.isin(labels_list, train_categories)
# get a list of indices to be used as train samples
train_indices = np.argwhere(labels_mask).squeeze()
# others will be used as test samples
test_indices = np.argwhere(np.logical_not(labels_mask)).squeeze()
# now that we have distinct indices for train and test sets, we can use `Subset` to create new datasets
# from `full_dataset`, which contain only the samples at given indices.
# finally, we apply transformations created above.
train_dataset = CarsDataset(
Subset(full_dataset, train_indices), transform=transform
)
test_dataset = CarsDataset(
Subset(full_dataset, test_indices), transform=transform
)
return train_dataset, test_dataset
def get_dataloaders(
batch_size: int,
input_size: int,
shuffle: bool = False,
):
train_dataset, test_dataset = get_datasets(input_size)
train_dataloader = GroupSimilarityDataLoader(
train_dataset, batch_size=batch_size, shuffle=shuffle
)
test_dataloader = GroupSimilarityDataLoader(
test_dataset, batch_size=batch_size, shuffle=False
)
return train_dataloader, test_dataloader
class CarsDataset(Dataset):
def __init__(self, dataset: Dataset, transform: Callable):
self._dataset = dataset
self._transform = transform
def __len__(self) -> int:
return len(self._dataset)
def __getitem__(self, index) -> SimilarityGroupSample:
image, label = self._dataset[index]
image = self._transform(image)
return SimilarityGroupSample(obj=image, group=label)
可训练模型
现在是时候回顾 Quaterion 最令人兴奋的构建块之一了：TrainableModel。它是您希望配置进行训练的模型的基类，它提供了多个以 configure_ 开头的钩子方法，用于设置训练阶段的各个方面，就像其自身的基类 pl.LightningModule 一样。它对于使用 Quaterion 进行微调至关重要，因此我们将分解 models.py 中的这段核心代码，并单独审查每个方法。让我们从导入开始：
import torch
import torchvision
from quaterion_models.encoders import Encoder
from quaterion_models.heads import EncoderHead, SkipConnectionHead
from torch import nn
from typing import Dict, Union, Optional, List
from quaterion import TrainableModel
from quaterion.eval.attached_metric import AttachedMetric
from quaterion.eval.group import RetrievalRPrecision
from quaterion.loss import SimilarityLoss, TripletLoss
from quaterion.train.cache import CacheConfig, CacheType
from .encoders import CarsEncoder
在以下代码片段中，我们继承 TrainableModel。您可以使用 __init__() 来存储一些属性，以便在后续各种 configure_* 方法中使用。然而，更有趣的部分在于 configure_encoders() 方法。我们需要从此方法返回一个 Encoder 实例（或一个以 Encoder 实例为值的字典）。在我们的例子中，它是一个 CarsEncoders 实例，我们很快会详细介绍它。现在请注意它是如何使用一个预训练的 ResNet152 模型创建的，该模型的分类层被替换为恒等函数。
class Model(TrainableModel):
def __init__(self, lr: float, mining: str):
self._lr = lr
self._mining = mining
super().__init__()
def configure_encoders(self) -> Union[Encoder, Dict[str, Encoder]]:
pre_trained_encoder = torchvision.models.resnet152(pretrained=True)
pre_trained_encoder.fc = nn.Identity()
return CarsEncoder(pre_trained_encoder)
在 Quaterion 中，一个 SimilarityModel 由一个或多个 Encoder 和一个 EncoderHead 组成。quaterion_models 提供了 多种 EncoderHead 实现，具有统一的 API，例如可配置的 dropout 值。您可以使用其中一种，也可以创建自己的 EncoderHead 子类。无论哪种情况，您都需要从 configure_head 方法返回一个 EncoderHead 实例。在此示例中，我们将使用一个 SkipConnectionHead，它轻量且更抗过拟合。
def configure_head(self, input_embedding_size) -> EncoderHead:
return SkipConnectionHead(input_embedding_size, dropout=0.1)
Quaterion 实现了 一些流行的相似度学习损失函数，所有这些都继承自 GroupLoss 或 PairwiseLoss。在此示例中，我们将使用 TripletLoss，它是 GroupLoss 的子类。一般来说，GroupLoss 的子类用于样本被分配到某个组（或标签）的数据集。在我们的示例中，标签是汽车的品牌。这些数据集应该生成 SimilarityGroupSample。其他替代方案是 PairwiseLoss 的实现，它处理 SimilarityPairSample - 一对单独指定相似度的对象。要查看后者的示例，您可能需要查阅 NLP 教程。
def configure_loss(self) -> SimilarityLoss:
return TripletLoss(mining=self._mining, margin=0.5)
configure_optimizers() 对于 PyTorch Lightning 用户来说可能很熟悉，但该方法内部使用了一个新的 self.model。它是 SimilarityModel 的一个实例，由 Quaterion 根据 configure_encoders() 和 configure_head() 的返回值自动创建。
def configure_optimizers(self):
optimizer = torch.optim.Adam(self.model.parameters(), self._lr)
return optimizer
Quaterion 中的缓存用于避免在每个 epoch 中重复计算冻结的预训练 Encoder 的输出。配置后，输出将计算一次并缓存在首选设备中，供后续直接使用。这既提供了显著的加速，也减少了内存占用。然而，它也相当灵活，有几个旋钮可以调整。为了充分发挥其潜力，建议您查看 缓存教程。为了使本文内容完整，您需要从 configure_caches() 返回一个 CacheConfig 实例来指定与缓存相关的偏好设置，例如：
CacheType，即是将缓存存储在 CPU 还是 GPU 上，save_dir，即持久化缓存以供后续运行的路径，batch_size，即仅在创建缓存时使用的批量大小——实际训练中使用的批量大小可能不同。
def configure_caches(self) -> Optional[CacheConfig]:
return CacheConfig(
cache_type=CacheType.AUTO, save_dir="./cache_dir", batch_size=32
)
我们刚刚配置了 TrainableModel 的训练相关设置。然而，评估是机器学习实验中不可或缺的一部分，您可以通过从 configure_metrics() 返回一个或多个 AttachedMetric 实例来配置评估指标。Quaterion 内置了多种 group 和 pairwise 评估指标。
def configure_metrics(self) -> Union[AttachedMetric, List[AttachedMetric]]:
return AttachedMetric(
"rrp",
metric=RetrievalRPrecision(),
prog_bar=True,
on_epoch=True,
on_step=False,
)
编码器
如前所述，一个 SimilarityModel 由一个或多个 Encoder 和一个 EncoderHead 组成。即使我们冻结了预训练的 Encoder 实例，EncoderHead 仍然是可训练的，并且具有足够的参数来适应手头的新任务。建议您尽可能将 trainable 属性设置为 False，因为这可以让您受益于上述缓存机制。另一个重要属性是 embedding_size，它将作为 input_embedding_size 传递给 TrainableModel.configure_head()，以便您正确初始化头部层。让我们看看以下借用自 encoders.py 的代码中如何实现一个 Encoder：
import os
import torch
import torch.nn as nn
from quaterion_models.encoders import Encoder
class CarsEncoder(Encoder):
def __init__(self, encoder_model: nn.Module):
super().__init__()
self._encoder = encoder_model
self._embedding_size = 2048 # last dimension from the ResNet model
@property
def trainable(self) -> bool:
return False
@property
def embedding_size(self) -> int:
return self._embedding_size
Encoder 是一个常规的 torch.nn.Module 子类，我们需要在 forward 方法中实现前向传播逻辑。根据您创建子模块的方式，此方法可能会更复杂；但是，在此示例中，我们仅通过预训练的 ResNet152 主干网络传递输入。
def forward(self, images):
embeddings = self._encoder.forward(images)
return embeddings
机器学习开发的一个重要步骤是正确地保存和加载模型。Quaterion 允许您使用 TrainableModel.save_servable() 保存 SimilarityModel，并使用 SimilarityModel.load() 恢复它。为了能够使用这两个方法，您需要在您的 Encoder 中实现 save() 和 load() 方法。此外，同样重要的是，您应将 Encoder 的子类定义在 __main__ 命名空间之外，即与您的主要入口文件分开存放。否则可能无法正确恢复。
def save(self, output_path: str):
os.makedirs(output_path, exist_ok=True)
torch.save(self._encoder, os.path.join(output_path, "encoder.pth"))
@classmethod
def load(cls, input_path):
encoder_model = torch.load(os.path.join(input_path, "encoder.pth"))
return CarsEncoder(encoder_model)
训练
在实现所有必需对象后，可以轻松地将它们组合在一起，并使用 Quaterion.fit() 方法运行训练循环。它需要：
- 一个
TrainableModel， - 一个
pl.Trainer， - 用于训练数据的
SimilarityDataLoader， - 以及可选的另一个用于评估数据的
SimilarityDataLoader。
我们需要导入一些对象来准备所有这些：
import os
import pytorch_lightning as pl
import torch
from pytorch_lightning.callbacks import EarlyStopping, ModelSummary
from quaterion import Quaterion
from .data import get_dataloaders
from .models import Model
以下代码片段中的 train() 函数需要几个超参数值作为参数。这些参数可以在 config.py 中定义或从命令行传递。然而，为了简洁起见，这部分代码被省略了。相反，让我们重点关注所有构建块如何初始化并传递给负责运行整个循环的 Quaterion.fit() 方法。训练循环完成后，您只需调用 TrainableModel.save_servable() 来保存 SimilarityModel 实例的当前状态。
def train(
lr: float,
mining: str,
batch_size: int,
epochs: int,
input_size: int,
shuffle: bool,
save_dir: str,
):
model = Model(
lr=lr,
mining=mining,
)
train_dataloader, val_dataloader = get_dataloaders(
batch_size=batch_size, input_size=input_size, shuffle=shuffle
)
early_stopping = EarlyStopping(
monitor="validation_loss",
patience=50,
)
trainer = pl.Trainer(
gpus=1 if torch.cuda.is_available() else 0,
max_epochs=epochs,
callbacks=[early_stopping, ModelSummary(max_depth=3)],
enable_checkpointing=False,
log_every_n_steps=1,
)
Quaterion.fit(
trainable_model=model,
trainer=trainer,
train_dataloader=train_dataloader,
val_dataloader=val_dataloader,
)
model.save_servable(save_dir)
评估
让我们看看通过这些简单步骤取得了什么成果。evaluate.py 中有两个函数用于评估基准模型和调优后的相似度模型。为了简洁起见，我们只介绍后者。除了轻松恢复 SimilarityModel 外，此代码片段还展示了如何使用 Evaluator 通过给定的评估指标来评估 SimilarityModel 在给定数据集上的性能。
原始模型与调优模型在检索上的比较
数据集的完整评估通常会呈指数增长，因此您可能希望对抽样的子集进行部分评估。在这种情况下，您可以使用 采样器 来限制评估范围。与用于训练的 Quaterion.fit() 类似，Quaterion.evaluate() 运行完整的评估循环。它接受以下参数：
- 一个使用给定评估指标和
Sampler创建的Evaluator实例， - 要评估的
SimilarityModel， - 以及评估数据集。
def eval_tuned_encoder(dataset, device):
print("Evaluating tuned encoder...")
tuned_cars_model = SimilarityModel.load(
os.path.join(os.path.dirname(__file__), "cars_encoders")
).to(device)
tuned_cars_model.eval()
result = Quaterion.evaluate(
evaluator=Evaluator(
metrics=RetrievalRPrecision(),
sampler=GroupSampler(sample_size=1000, device=device, log_progress=True),
),
model=tuned_cars_model,
dataset=dataset,
)
print(result)
结论
在本教程中，我们训练了一个相似度模型，用于从训练阶段未见过的新类别中搜索相似的汽车。然后，我们使用检索 R-Precision 指标在测试数据集上对其进行了评估。基础模型得分 0.1207，而我们的调优模型达到 0.2540，是其两倍多的分数。这些分数如下图所示：
基础模型和调优模型的指标