Векторизация вычислений VLAD в numpy

Мне было интересно, можно ли векторизовать эту реализацию вычислений VLAD.

Для контекста:

feats = numpy массив формы (T, N, F)

kmeans = Объект KMeans из scikit-learn инициализирован с помощью K кластеры.

Текущий метод

      k = kmeans.n_clusters # K
centers = kmeans.cluster_centers_ # (K, F)
vlad_feats = []

for feat in feats:
    # feat shape - (N, F) 
    cluster_label = kmeans.predict(feat) #(N,)
    vlad = np.zeros((k, feat.shape[1])) # (K, F)

    # computing the differences for all the clusters (visual words)
    for i in range(k):
        # if there is at least one descriptor in that cluster
        if np.sum(cluster_label == i) > 0:
            # add the differences
            vlad[i] = np.sum(feat[cluster_label == i, :] - centers[i], axis=0)
    vlad = vlad.flatten() # (K*F,)
    # L2 normalization
    vlad = vlad / np.sqrt(np.dot(vlad, vlad))
    vlad_feats.append(vlad)

vlad_feats = np.array(vlad_feats) # (T, K*F)

Получение прогнозов kmeans в виде пакета не проблема, поскольку мы можем сделать следующее:

      feats2 = feats.reshape(-1, F) # (T*N, F)
labels = kmeans.predict(feats2) # (T*N, )

Но я застрял в вычислении расстояний между кластерами.

2 ответа

Решение

Вы начали с правильного подхода. Попробуем вытащить все линии из петли одну за другой. Во-первых, прогнозы:

      cluster_label = kmeans.predict(feats.reshape(-1, F)).reshape(T, N)  # T, N

Вам действительно не нужен чек np.sum(cluster_label == i) > 0, так как сумма все равно окажется нулевой. Ваша цель - сложить расстояния от центра для каждой метки в каждом T и особенность.

Вы можете вычислить маски cluster_label == iиспользуя простую трансляцию. Вы хотите, чтобы последнее измерение было K:

      mask = cluster_label[..., None] == np.arange(k)   # T, N, K

Вы также можете вычислить k различия feats - centers[i] используя более сложную трансляцию:

      delta = feats[..., None, :] - centers # T, N, K, F

Теперь вы можете умножить различия на маску и уменьшить по N размер путем суммирования:

      vlad = (delta * mask[..., None]).sum(axis=1).reshape(T, -1)  # T, K * F

Отсюда нормализация должна быть тривиальной:

      vlad /= np.linalg.norm(vlad, axis=1, keepdims=True)

Хотя ответ @MadPhysicist векторизуется, я обнаружил, что это вредит производительности.

Ниже, looping по сути является переписанной версией алгоритма OP и naivec использует векторизацию с помощью разнесенного тензора 4D.

      import numpy as np
from sklearn.cluster import MiniBatchKMeans

def looping(kmeans: MiniBatchKMeans, local_tlf):
    k, (t, l, f) = kmeans.n_clusters, local_tlf.shape
    centers_kf = kmeans.cluster_centers_
    vlad_tkf = np.zeros((t, k, f))
    for vlad_kf, local_lf in zip(vlad_tkf, local_tlf):
        label_l = kmeans.predict(local_lf)
        for i in range(k):
            vlad_kf[i] = np.sum(local_lf[label_l == i] - centers_kf[i], axis=0)
        vlad_D = vlad_kf.ravel()
        vlad_D = np.sign(vlad_D) * np.sqrt(np.abs(vlad_D))
        vlad_D /= np.linalg.norm(vlad_D)
        vlad_kf[:,:] = vlad_D.reshape(k, f)
    return vlad_tkf.reshape(t, -1)


def naivec(kmeans: MiniBatchKMeans, local_tlf):
    k, (t, l, f) = kmeans.n_clusters, local_tlf.shape
    centers_kf = kmeans.cluster_centers_
    labels_tl = kmeans.predict(local_tlf.reshape(-1,f)).reshape(t, l)
    mask_tlk = labels_tl[..., np.newaxis] == np.arange(k)
    local_tl1f = local_tlf[...,np.newaxis,:]
    delta_tlkf = local_tl1f - centers_kf # <-- easy to run out of memory
    vlad_tD = (delta_tlkf * mask_tlk[..., np.newaxis]).sum(axis=1).reshape(t, -1)
    vlad_tD = np.sign(vlad_tD) * np.sqrt(np.abs(vlad_tD))
    vlad_tD /= np.linalg.norm(vlad_tD, axis=1, keepdims=True)
    return vlad_tD

В самом деле, см. Пример ниже.

      np.random.seed(1234)
# usually there are a lot more images than this
t, l, f, k = 256, 128, 64, 512
X = np.random.randn(t, l, f)
km = MiniBatchKMeans(n_clusters=16, n_init=10, random_state=0)
km.fit(X.reshape(-1, f))

result_looping = looping(km, X)
result_naivec = naivec(km, X)

%timeit looping(km, X) # ~200 ms
%timeit naivec(km, X) # ~300 ms

assert np.allclose(result_looping, result_naivec)

Идиоматическая векторизация, которая предотвращает рост памяти за пределы выходного размера (асимптотически), будет способствовать уменьшению разброса.

      def truvec(kmeans: MiniBatchKMeans, local_tlf):
    k, (t, l, f) = kmeans.n_clusters, local_tlf.shape
    centers_kf = kmeans.cluster_centers_
    labels_tl = kmeans.predict(local_tlf.reshape(-1,f)).reshape(t, l)
    
    vlad_tkf = np.zeros((t, k, f))
    M = t * k
    labels_tl += np.arange(t)[:, np.newaxis] * k
    vlad_Mf = vlad_tkf.reshape(-1, f)
    np.add.at(vlad_Mf, labels_tl.ravel(), local_tlf.reshape(-1, f))
    counts_M = np.bincount(labels_tl.ravel(), minlength=M)
    vlad_tkf -= counts_M.reshape(t, k, 1) * centers_kf
    
    vlad_tD = vlad_tkf.reshape(t, -1)
    vlad_tD = np.sign(vlad_tD) * np.sqrt(np.abs(vlad_tD))
    vlad_tD /= np.linalg.norm(vlad_tD, axis=1, keepdims=True)
    return vlad_tD

Однако, к сожалению, это также дает нам понять 200 msвремя вычисления. Это сводится к двум причинам:

  • внутренний цикл уже векторизован в
  • np.add.atфактически не может использовать векторизованные инструкции ЦП, в отличие от оригинального поэтапного сокращения np.sum(local_lf[label_l == i] - centers_kf[i], axis=0)

Производительная векторизованная версия алгоритма VLAD требует некоторых сложных методов для использования непрерывного доступа к массивам. Эта версия улучшена на 40% по сравнению с looping(), но требует большой настройки --- см. мой блог о подходе здесь.

Другие вопросы по тегам