tensorflow實現kNN contrastive loss

Notes

題目所謂 kNN contrastive loss 是我自己起的名字，指的是 [1] 的第 2 條損失，文中叫做 Semi-supervised Embedding Term，形式如下：
$J(x_i, x_j)=\begin{cases} d(x_i, x_j), & A(i,j)=1 \\ \max\{0,m-d(x_i, x_j)\}, & A(i,j)=0 \end{cases}$
其中 $d(x_i,x_j)$ 表示一種距離（文中是歐氏距離平方），m 是 margin，
$A(i,j)=\begin{cases} 1, & U(i,j)\wedge x_j\in NN_k(i) \\ 0, & U(i,j)\wedge x_j\notin NN_k(i) \end{cases}$
其中 $U(i,j)$ 表示 $x_i$ 和 $x_j$ 都是 unlabeled 的 至少有一個是 unlabeled 的，$x_j\in NN_k(i)$ 表示 $x_j$（的 embedding）在 $x_i$（的 embedding）的 k 近鄰之內。
所以這條 loss 是半監督的正則項，對於抽樣的樣本對 $(x_i,x_j)$，如果在 kNN 內，就希望拉近兩者；否則希望將兩者的距離拉大到至少 m。
文中說爲了正負例的平衡，對於一個 batch 內的每個 sample，都在 batch 內最近的 k 個裏抽一個做 positive sample （k 還和 $\frac{batch\ size}{3}$ 取了 min，是個 trick 吧）、最遠的幾個裏抽一個做 negative sample （batch size 是 128，它選 neg 的範圍是排最後的 [120, 124)，也是 trick 吧） 。

Mask

這裏實現一個初級的版本：最近的 k_pos 個裏隨機選一個、最遠的 k_neg 個裏隨機選一個。
大思路是靠 mask，參考 [3]，用到tf.scatter_nd，也用到 [5] 的隨機索引。
效果是傳入距離矩陣 D，返回一個同形的矩陣 M，$M_{i,j}=1\Leftrightarrow 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索引