RexKing6's Note

三元组损失与TensorFlow在线三元组挖掘

发表于 更新于 分类于 ， ， 阅读次数：
12k 11 分钟

简介

本文是翻译Triplet Loss and Online Triplet Mining in TensorFlow而来。

第一次尝试独自翻译外文博客，不通顺的地方多见谅。

原作者写这个博客一开始是因为有人在stackflow提了一个问题问，“TensorFlow怎么实现三元组损失？”原作者便在下面给出了回答，“计算三元组并不困难，困难在于怎么挖掘得到三元组”，提问者问了原作者是否有博客能提供参考。时隔两年，在前两个月，原作者发布了这篇博客。

接下来，开始正文，正文中的“我”是指原作者。

引言

在人脸识别中， 三元组损失被用来学习好的脸部嵌入矩阵（或“编码矩阵”）。如果你对三元组损失不太熟悉，你应该先从吴恩达的深度学习课程了解学习一下。

众所周知，三元组损失是难以实现的，特别是在添加了使用TensorFlow构建计算图这个约束。

在这篇博客中，我将定义三元组损失和取样得到三元组的不同策略。然后，我将解释如何正确地在TensorFlow中实现在线挖掘得到三元组损失。

大约在两年前，我在Reminiz实习，做人脸识别。我在stackoverflow上回答了一个关于在TensorFlow中实现三元组损失的问题。我得出的结论：

显然，在Tensorflow中实现三元组损失是很困难的，而且有很多方法可以使它比在python中进行采样更有效，但是解释它们需要一个完整的博客！

两年后了，我们开始吧。

所有的代码都可以在这个github库中找到。

三元组损失与三元组在线挖掘

为什么不使用softmax？

Google的一篇论文“FaceNet: A Unified Embedding for Face Recognition and Clustering”介绍了人脸识别的三元组损失。他们描述了一种使用在线挖掘三元组来训练脸部嵌入矩阵的新方法，这将在下一节中讨论。

通常在监督学习中，我们有固定数量的类，并使用软-最大交叉熵损失来训练网络。然而，在某些情况下，我们需要能够有一个可变数量的类。例如，在人脸识别中，我们需要能够比较两个未知的面孔，并判断它们是否来自同一个人。

在这种情况下，三元组损失是一种学习对每个脸部的良好嵌入的方法。在嵌入空间中，来自同一个人的面孔应该紧密地靠近在一起，形成良好的分离的集群。

损失定义


两张同类脸部（Obama）和一张异类脸部（Macron）

三元组损失的目标是确保：

  • 两个具有相同标签的样本在嵌入空间中它们的嵌入紧密靠近；
  • 两个具有不同标签的样本在嵌入空间中它们的嵌入距离很远。

但是，我们不想把每个样本的训练嵌入推移缩小到非常小的集群。唯一要求是给定两个正（同类）样本和一个负（异类）样本。负的距离应该要比正的距离更大一些。这与SVM中使用的间隔非常相似，这里我们希望每个类的集群都由该间隔分隔开。

为了正式定义该要求，损失函数由嵌入矩阵的三元组定义而成：

  • 一个anchor
  • 一个与anchor同类的正样本
  • 一个与anchor异类的负样本

用于嵌入空间的距离度量$d$，三元$(a,p,n)$的损失为：

我们最小化损失，它使$d(a, p)$趋向于0，使$d(a, n)$大于$d(a, p) + margin$。一旦$n$为简单的负样本，损失即为0。

三元组挖掘

根据损失的定义，由三种类型的三元组：

  • 简单三元组：损失为0的三元组，因为$d(a, p) + margin < d(a,n)$；
  • 困难三元组：在三元组中，负样本比正样本更靠近anchor，即$d(a,n) < d(a,p)$；
  • 半困难三元组：三元组中，负样本比正样本更不靠近anchor，但仍然有大于0的损失：$d(a, p) < d(a, n) < d(a, p) + margin$

每一个定义都依赖于负样本与anchor和正样本的相对位置。因此，我们可以将这三种类型扩展到负样本：困难负样本、半困难负样本或简单负样本。

下图显示了负样本的嵌入空间的三个对应区域。


给定一个anchor和一个正样本，负样本的三种情况

选择我们想要训练的三元组将会极大地影响我们的指标。在最初的Facenet论文，他们为每对anchor与正样本选取一个随机的半困难负样本，并且在这些三元组上训练。

离线与在线三元组挖掘

我们已经定义了一个三元组的损失，并且已经看到一些三元组比其他的更有用。现在的问题是如何对这些三元组进行取样或“挖掘”。

离线三元组挖掘

第一种产生三元组的方式是在每个epoch开始之前，离线找到它们。我们计算训练集的所有嵌入，然后只选择困难或半困难的三元组。然后我们就可以在这些三元组上训练一个epoch。

具体而言，我们会产生一个三元组$(i, j, k)$的列表。我们会构建$B$个三元组的批次，这意味着我们将计算$3B$个嵌入得到$B$个三元组，之后计算这些三元组的损失并且在网络中反向传播。

总的来说，这种方式不是很有效，因为我们需要对训练集进行全面遍历来生成三元组。它还需要定期离线地更新三元组。

在线三元组挖掘

在Facenet中引入在线三元组挖掘，这在Brandon Amos的博客“OpenFace 0.2.0: Higher accuracy and halved execution time”得到很好的描述。

这里的想法是，为每一批输入计算出有用的三元组。给定一个$B$个样本的批次（即有脸部的$B$张图像），我们计算$B$的嵌入，然后我们可以找到最大为$B^3$个三元组，其中大多数三元组都是无效的（它们没有两个正样本和一个负样本）。

这种方式可以给你的一个batch输入提供更多三元组，并且不需要任何离线挖掘。因此它的效率更高。我们将在最后一部分看到这方面的实现。


在线挖掘的三元组损失：从一组嵌入中计算出三元组。

在线挖掘的策略

在线挖掘中，我们已经从一个$B$输入的批次中计算该批次$B$的嵌入。现在我们想从这些$B$嵌入中生成三元组。

我们有三个索引$i,j,k \in [1,B]$，如果样本$i$和$j$有相同的标签但却相距较远，$k$样本有不同的标签，我们说$(i,j,k)$是一个有效的三元组。剩下的就是要有一个好的策略，在有效的三元组中挑选出一些来计算损失。

下面这两种策略的详细解释可以在这篇论文In Defense of the Triplet Loss for Person Re-Identification的第二部分找到。

他们假设你有一组面孔作为$B=PK$大小的输入，由P个不同的人，每个人有K张图像。一个经典的值是$K=4$。两种策略是：

  • batch all：选择所有有效的三胞胎，并在困难的和半困难的三胞胎中平均损失。
    • 这里的关键点是不考虑简单的三元组（即那些损失为0的），因为把它们加入平均后会使整体的损失非常小；
    • 这会产生总共有$PK(K-1)(PK-K)$个三元组（$PK$个anchor，每个anchor有$K-1$个可能的正样本，$PK-K$可能的负样本）。
  • batch hard：对于每个anchor，在批次中选取出最困难的正样本（$d(a,p)$最大）和最困难的负样本
    • 这会产生$PK$个三元组；
    • 被选取的三元组是批次中最困难的。

根据上面提到的论文，batch hard获得了最佳表现：

此外，所选的三元组可以被认为是适合的三元组，因为它们在数据的一小部分中是最困难的，这正是最适合学习三元组损失的方法。

然而，这实际上取决于你的数据集，应该通过比较开发的数据集的表现来决定。

一种简单的三元组损失的实现

在stackoverflow的回答中，我简单地实现了离线三元组挖掘的三元组损失：

1
2
3
4
5
6
7
8
9
anchor_output = ...    # shape [None, 128]
positive_output = ...  # shape [None, 128]
negative_output = ...  # shape [None, 128]

d_pos = tf.reduce_sum(tf.square(anchor_output - positive_output), 1)
d_neg = tf.reduce_sum(tf.square(anchor_output - negative_output), 1)

loss = tf.maximum(0.0, margin + d_pos - d_neg)
loss = tf.reduce_mean(loss)

该网络将被复制三次（用共享权重）来产生$B$个anchor，$B$个正样本，$B$个负样本的嵌入。然后我们只需计算这些嵌入的三元组损失。

这是一个简单的实现，但也是一个非常低效的实现，因为它使用的是离线的三元组挖掘。

一个更好的在线三元组挖掘的实现

所有相关的代码都可以在github上的model/triplet_loss.py有提供。

TensorFlow本身有一种半困难的在线挖掘三元组损失的实现，tf.contrib.losses.metric_learning.triplet_semihard_loss。在这里，我们不会遵循这个实现，并从头开始。

计算距离矩阵

因为最后的三元组损失取决于$d(a,p)$和$d(a,n)$，我们首先需要有效地计算成对的距离矩阵。我们在_pairwise_distances实现了欧式范数与欧式范数的平方。

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
def _pairwise_distances(embeddings, squared=False):
    """Compute the 2D matrix of distances between all the embeddings.

    Args:
        embeddings: tensor of shape (batch_size, embed_dim)
        squared: Boolean. If true, output is the pairwise squared euclidean distance matrix.
                 If false, output is the pairwise euclidean distance matrix.

    Returns:
        pairwise_distances: tensor of shape (batch_size, batch_size)
    """
    # Get the dot product between all embeddings
    # shape (batch_size, batch_size)
    dot_product = tf.matmul(embeddings, tf.transpose(embeddings))

    # Get squared L2 norm for each embedding. We can just take the diagonal of `dot_product`.
    # This also provides more numerical stability (the diagonal of the result will be exactly 0).
    # shape (batch_size,)
    square_norm = tf.diag_part(dot_product)

    # Compute the pairwise distance matrix as we have:
    # ||a - b||^2 = ||a||^2  - 2 <a, b> + ||b||^2
    # shape (batch_size, batch_size)
    distances = tf.expand_dims(square_norm, 0) - 2.0 * dot_product + tf.expand_dims(square_norm, 1)

    # Because of computation errors, some distances might be negative so we put everything >= 0.0
    distances = tf.maximum(distances, 0.0)

    if not squared:
        # Because the gradient of sqrt is infinite when distances == 0.0 (ex: on the diagonal)
        # we need to add a small epsilon where distances == 0.0
        mask = tf.to_float(tf.equal(distances, 0.0))
        distances = distances + mask * 1e-16

        distances = tf.sqrt(distances)

        # Correct the epsilon added: set the distances on the mask to be exactly 0.0
        distances = distances * (1.0 - mask)

    return distances

为了更详细地解释代码，我们计算嵌入的点积，维数为$(B,B)$。每个嵌入的欧式范数平方实际上就是在这个点积的对角线上，所以我们用tf.diag_part把它提取出来。最后我们用公式来计算距离：

一个棘手的问题是如果squared=False，我们取距离矩阵的平方根。首先，我们要确保距离矩阵总是正的。有些值可能是负的，因为计算中的小误差。我们只需要确保每个负值都被设置为0.0。

第二件要注意的事情是如果任何元素都是0.0的话（例如对角线应该都是0.0）。因为$0$的平方根导数是无限的，我们将得到的梯度为nan。为了处理这种情况，我们把等于0.0的值替换为一个比较小的值epsilon=1e-16。然后我们计算平方根，替换$\sqrt{\epsilon}$为0.0的值。

batch all 策略

在这个策略中，我们想要计算几乎所有三元组的三元组损失。在TensorFlow的计算图中，我们构造了3维的Tensor，维数为$(B,B,B)$，其中索引$(i,j,k)$，包含着它们的三元组损失。

然后我们得到一个有效三元组的3维mask_get_triplet_mask。这里mask[i,j,k]能正确识别$(i,j,k)$是一个有效的三元组。

最后，我们把无效的三元组的损失设置为0，并且计算有效的三元组损失的平均值。

这一切都是在batch_all_triplet_loss函数中实现的。

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
def batch_all_triplet_loss(labels, embeddings, margin, squared=False):
    """Build the triplet loss over a batch of embeddings.

    We generate all the valid triplets and average the loss over the positive ones.

    Args:
        labels: labels of the batch, of size (batch_size,)
        embeddings: tensor of shape (batch_size, embed_dim)
        margin: margin for triplet loss
        squared: Boolean. If true, output is the pairwise squared euclidean distance matrix.
                 If false, output is the pairwise euclidean distance matrix.

    Returns:
        triplet_loss: scalar tensor containing the triplet loss
    """
    # Get the pairwise distance matrix
    pairwise_dist = _pairwise_distances(embeddings, squared=squared)

    anchor_positive_dist = tf.expand_dims(pairwise_dist, 2)
    anchor_negative_dist = tf.expand_dims(pairwise_dist, 1)

    # Compute a 3D tensor of size (batch_size, batch_size, batch_size)
    # triplet_loss[i, j, k] will contain the triplet loss of anchor=i, positive=j, negative=k
    # Uses broadcasting where the 1st argument has shape (batch_size, batch_size, 1)
    # and the 2nd (batch_size, 1, batch_size)
    triplet_loss = anchor_positive_dist - anchor_negative_dist + margin

    # Put to zero the invalid triplets
    # (where label(a) != label(p) or label(n)  == label(a) or a == p)
    mask = _get_triplet_mask(labels)
    mask = tf.to_float(mask)
    triplet_loss = tf.multiply(mask, triplet_loss)

    # Remove negative losses (i.e. the easy triplets)
    triplet_loss = tf.maximum(triplet_loss, 0.0)

    # Count number of positive triplets (where triplet_loss > 0)
    valid_triplets = tf.to_float(tf.greater(triplet_loss, 1e-16))
    num_positive_triplets = tf.reduce_sum(valid_triplets)
    num_valid_triplets = tf.reduce_sum(mask)
    fraction_positive_triplets = num_positive_triplets / (num_valid_triplets + 1e-16)

    # Get final mean triplet loss over the positive valid triplets
    triplet_loss = tf.reduce_sum(triplet_loss) / (num_positive_triplets + 1e-16)

    return triplet_loss, fraction_positive_triplets

_get_triplet_mask的实现非常简单，所以我不会详细说明。

batch hard 策略

在这个策略中，我们想要找到对每个anchor最困难的正样本和负样本。

最困难的正样本

为了计算最困难的正样本，我们从成对的距离矩阵开始。然后我们得到一个有效配对的二维mask$(a,p)$（即$a \neq p$并且$a$和$p$有相同的标签），之后把mask外的其他元素都置为0。

最后一步是取这个经过修改的距离矩阵的每一行的最大距离。结果应该是有效的对$(a,p)$，因为不有效的元素都置为0。

最困难的负样本

最困难的负样本也是类似的，但要计算起来有点难。这里我们需要得到每一行的最小距离，所以我们不能把不有效的对$(a,n)$置为0（不有效是指$a$和$n$有相同的标签）。

这里的技巧是每一行都要为无效的对$(a,n)$添加最大的值。之后然后我们用最小值除以每一行。结果应该是有效的对$(a,n)$，因为不有效的元素都置为最大值。

最后一步是将这些合并到三元组损失中：

1
triplet_loss = tf.maximum(hardest_positive_dist - hardest_negative_dist + margin, 0.0)

这一切都是在batch_hard_triplet_loss函数中实现的。

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
def batch_hard_triplet_loss(labels, embeddings, margin, squared=False):
    """Build the triplet loss over a batch of embeddings.

    For each anchor, we get the hardest positive and hardest negative to form a triplet.

    Args:
        labels: labels of the batch, of size (batch_size,)
        embeddings: tensor of shape (batch_size, embed_dim)
        margin: margin for triplet loss
        squared: Boolean. If true, output is the pairwise squared euclidean distance matrix.
                 If false, output is the pairwise euclidean distance matrix.

    Returns:
        triplet_loss: scalar tensor containing the triplet loss
    """
    # Get the pairwise distance matrix
    pairwise_dist = _pairwise_distances(embeddings, squared=squared)

    # For each anchor, get the hardest positive
    # First, we need to get a mask for every valid positive (they should have same label)
    mask_anchor_positive = _get_anchor_positive_triplet_mask(labels)
    mask_anchor_positive = tf.to_float(mask_anchor_positive)

    # We put to 0 any element where (a, p) is not valid (valid if a != p and label(a) == label(p))
    anchor_positive_dist = tf.multiply(mask_anchor_positive, pairwise_dist)

    # shape (batch_size, 1)
    hardest_positive_dist = tf.reduce_max(anchor_positive_dist, axis=1, keepdims=True)

    # For each anchor, get the hardest negative
    # First, we need to get a mask for every valid negative (they should have different labels)
    mask_anchor_negative = _get_anchor_negative_triplet_mask(labels)
    mask_anchor_negative = tf.to_float(mask_anchor_negative)

    # We add the maximum value in each row to the invalid negatives (label(a) == label(n))
    max_anchor_negative_dist = tf.reduce_max(pairwise_dist, axis=1, keepdims=True)
    anchor_negative_dist = pairwise_dist + max_anchor_negative_dist * (1.0 - mask_anchor_negative)

    # shape (batch_size,)
    hardest_negative_dist = tf.reduce_min(anchor_negative_dist, axis=1, keepdims=True)

    # Combine biggest d(a, p) and smallest d(a, n) into final triplet loss
    triplet_loss = tf.maximum(hardest_positive_dist - hardest_negative_dist + margin, 0.0)

    # Get final mean triplet loss
    triplet_loss = tf.reduce_mean(triplet_loss)

    return triplet_loss

测试我们的实现

如果你不相信上面的实现是按预期工作的，那么你是对的！要确保实现中没有bug，唯一的方法就是为model/triplet_loss.py的每个函数编写测试。

这对于像这样的棘手函数来说尤其重要，因为在TensorFlow中很难实现，但是在python中使用三个嵌套for循环来编写更容易。

测试函数写在model/tests/test_triplet_loss.py，并且将我们的TensorFlow实现的结果与简单的numpy实现的结果进行比较。

为了检查测试是否通过，运行：

1
pytest model/tests/test_triplet_loss.py

（或只是pytest）

下面是执行的测试列表：

  • test_pairwise_distances()：比较了TensorFlow与numpy的成对距离的结果。
  • test_pairwise_distances_are_positive()：确保得到的距离是正的。
  • test_gradients_pairwise_distances()：确保梯度不是nan。
  • test_triplet_mask()：比较numpy和tensorflow实现。
  • test_anchor_positive_triplet_mask()：比较numpy和tensorflow实现。
  • test_anchor_negative_triplet_mask()：比较numpy和tensorflow实现。
  • test_simple_batch_all_triplet_loss()：只有一种标签的测试。
  • test_batch_all_triplet_loss()：对batch all策略的全部测试（与numpy相比）
  • test_batch_hard_triplet_loss()：对batch hard策略的全部测试（与numpy相比）

在MNIST上实验

即使在上面的测试中，也很容易忽视一些错误。例如，一开始我实现了成对的距离而不检查根号下的值是否大于$0$。我通过了所有测试，但是训练过程中的梯度是nan。因此我加入了test_gradients_pairwise_distances，修正了_pairwise_distances函数。

为了使事情简单，我们将测试MNIST的三元组损失。代码可以在这里找到。

为了训练和评价模型，运行：

1
python train.py --model_dir experiments/base_model

这将启动一个新的实验（即一个训练过程）称为base_model。模型目录（包含权重、总结……）在experiments/base_model。这里我们使用了一个json文件experiments/base_model/params.json，指定了模型中的所有超参数。这个在所有新的实验中都应该被创建。

一旦训练过程完成（或者在模型目录中保存了一些权重），我们可以使用TensorBoard来可视化这些嵌入。要做到这一点，运行:

1
python visualize_embeddings.py --model_dir experiments/base_model

在实验目录中运行TensorBoard：

1
tensorboard --logdir experiments/base_model


用T-SNE可视化的MNIST测试图像的嵌入。

这些嵌入是由配置文件experiments/base_model/params.json中指定的超参数运行的。看到哪些验证集图片被错误分类是很有趣的：它们中的大多数肯定也会被人类误解。

结论

TensorFlow并不能很容易地实现三元组损失，但是只要稍加努力，我们就可以通过在线挖掘来构建一个漂亮的三元组损失。

最棘手的部分是如何有效地计算嵌入的距离，以及如何过滤掉无效的/简单的三胞胎。

最后，如果你需要记住一件事：总是测试你的代码，特别是当它像三元组损失一样复杂时。

资源

相关文章
一分一毛，也是心意。