tensorflow实现kNN contrastive loss

Notes

题目所谓 kNN contrastive loss 是我自己起的名字，指的是 [1] 的第 2 条损失，文中叫做 Semi-supervised Embedding Term，形式如下：
$$J(x_i,x_j)=\{\begin{array}{ll}d(x_i,x_j),& A(i,j)=1\\ \max \{0,m-d(x_i,x_j)\},& A(i,j)=0\end{array}$$
其中 $d(x_i,x_j)$ 表示一种距离（文中是欧氏距离平方），m 是 margin，
$$A(i,j)=\{\begin{array}{ll}1,& U(i,j)\wedge x_j\in N{N}_k(i)\\ 0,& U(i,j)\wedge x_j\notin N{N}_k(i)\end{array}$$
其中 $U(i,j)$ 表示 $x_i$ 和 $x_j$ 都是 unlabeled 的 至少有一个是 unlabeled 的，$x_j\in N{N}_k(i)$ 表示 $x_j$（的 embedding）在 $x_i$（的 embedding）的 k 近邻之内。
所以这条 loss 是半监督的正则项，对于抽样的样本对 $(x_i,x_j)$，如果在 kNN 内，就希望拉近两者；否则希望将两者的距离拉大到至少 m。
文中说为了正负例的平衡，对于一个 batch 内的每个 sample，都在 batch 内最近的 k 个里抽一个做 positive sample （k 还和 $\frac{batchsize}{3}$ 取了 min，是个 trick 吧）、最远的几个里抽一个做 negative sample （batch size 是 128，它选 neg 的范围是排最后的 [120, 124)，也是 trick 吧，相当于没那么 easy 的 mining） 。

Mask

这里实现一个初级的版本：最近的 k_pos 个里随机选一个、最远的 k_neg 个里随机选一个。
大思路是靠 mask，参考 [3]，用到tf.scatter_nd，也用到 [5] 的随机索引。
效果是传入距离矩阵 D，返回一个同形的矩阵 M，$M_{i,j}=1\iff d(i,j)$ 是第 i 行最大的 k 个元素之一。

#import tensorflow as tf
def _top_k_mask(D, k, rand_pick=False):
    """M[i][j] = 1 <=> D[i][j] is oen of the BIGGEST k in i-th row
    Args:
        D: (n, n), distance matrix
        k: param `k` of kNN
        rand_pick: true or false
            - if `False`, 普通 top-K mask，全部 top-K 都选
            - if `True`, 每行的 top-K 中再随机选一个
    """
    n_row = tf.shape(D)[0]
    n_col = tf.shape(D)[1]

    k_val, k_idx = tf.math.top_k(D, k)
    if rand_pick:  # 只留 top-K 中随机一个
        c_idx = tf.random_uniform([n_row, 1],
                                  minval=0, maxval=k,
                                  dtype="int32")
        r_idx = tf.range(n_row, dtype="int32")[:, None]
        idx = tf.concat([r_idx, c_idx], axis=1)
        k_idx = tf.gather_nd(k_idx, idx)[:, None]

    idx_offset = (tf.range(n_row) * n_col)[:, None]
    k_idx_linear = k_idx + idx_offset
    k_idx_flat = tf.reshape(k_idx_linear, [-1, 1])

    updates = tf.ones_like(k_idx_flat[:, 0], "int8")
    mask = tf.scatter_nd(k_idx_flat, updates, [n_row * n_col])
    mask = tf.reshape(mask, [-1, n_col])
    mask = tf.cast(mask, "bool")

    return mask

Loss

由上面的 mask 实现。labeled 和 unlabeled 样本分两个 placeholder 输入，pairwise_dist分开无标签对无标签、无标签对有标签两种，调两次下面的knn_loss相加。

#import tensorflow as tf
def knn_loss(pairwise_dist, k_pos, k_neg, margin):
    mask_knn_pos = _top_k_mask(-1.0 * pairwise_dist, k_pos, rand_pick=True)  # 最远
    mask_knn_neg = _top_k_mask(pairwise_dist, k_neg, rand_pick=True)  # 最近

    mask_knn_pos = tf.cast(mask_knn_pos, "float32")
    mask_knn_neg = tf.cast(mask_knn_neg, "float32")

    dis_pos = pairwise_dist
    dis_neg = tf.maximum(0.0, margin - pairwise_dist)
    knn_loss = (dis_pos * mask_knn_pos + dis_neg * mask_knn_neg) / 2.0

    valid_item = tf.to_float(tf.greater(knn_loss, 1e-16))
    n_valid = tf.reduce_sum(valid_item)
    knn_loss = tf.reduce_sum(knn_loss) / (n_valid + 1e-16)
    return knn_loss


"""两个距离"""
# unlabeled v.s. unlabeled
dist_uu = _euclidean_dist(feature_u, feature_u)
# unlabeled v.s. labeled
dist_ul = _euclidean_dist(feature_u, feature_l)
"""knn loss"""
loss_knn_uu = losses.knn_loss(dist_uu, k_pos, k_neg, margin)
loss_knn_ul = losses.knn_loss(dist_ul, k_pos, k_neg, margin)
loss_knn = loss_knn_uu + loss_knn_ul

References

1. SSDH: Semi-Supervised Deep Hashing for Large Scale Image Retrieval
2. PKU-ICST-MIPL/SSDH_TCSVT2017
3. 使用带有top_k输出的scatter_nd？
4. tensorflow实现triplet loss
5. tensorflow用gather/gather_nd实现tensor索引