Как GPU-вычисления буквально спасли меня на работе. Пример на Python

Забавно, но перенесённой на Torch и немного оптимизированной версии решения от Physmatik уже требуется всего 0.5 секунды на GPU (1080Ti) и 10.5 секунд на CPU (TR 1920X) на обработку тех же 10^5 векторов. Учитывая возможность эффективного использования FP16 на карте V100, то на ней будет ещё быстрее.

Код на Torch

import torch

n = 10**5
k = 59
batch_size = 10**3
device = 'cuda:0' # 'cpu'
nn_number = 3

items = np.random.rand(n, k).astype(np.float32)
items = items / np.sqrt(np.sum(items**2, axis=1, keepdims=True))

def three_closest(X, indices):
    sample_n = indices.shape[0]
    index = torch.zeros((sample_n, nn_number), dtype=torch.int64, device=device)
    distances = X[indices, :] @ X.T
    ind_range = torch.arange(sample_n, device=device)
    distances[ind_range, indices] = 0
    for j in range(nn_number):
        max_index = torch.argmax(distances, dim=1)
        index[:, j] = max_index
        distances[ind_range, max_index] = 0
    return index

def distance_func(X):
    X = torch.from_numpy(X).to(device=device)
    all_items = torch.zeros((n, nn_number), dtype=torch.int64, device=device)
    for indices in torch.split(torch.arange(n, device=device), batch_size):        
        all_items[indices] = three_closest(X, indices)
    return all_items.cpu().numpy()

nearest_torch = distance_func(items)

Как GPU-вычисления буквально спасли меня на работе. Пример на Python

Считается именно 20 секунд, а не 0.5 секунд (на самом деле 0.5 мс будет выведено, перезапустил код локально). Автор посчитал затраченное время на асинхронный вызов.
Также хочется отметить, число колонок 59 — очень неудачное. Если добавить ещё 5 (например, нулевых) до 64, то выполняться на gpu будет быстрее примерно в два раза.

И, конечно, стоило бы воспользоваться готовой библиотекой для этой задачи. Например, при помощи faiss те же 10^5 векторов можно обработать за 25 секунд на TR 1920x (против 50 секунд для кода из статьи на 1080Ti). Результаты совпадают.

Код

import numpy as np
import faiss


n = 10**5
k = 59
batch_size = 18

items = np.random.rand(n, k).astype(np.float32)
items = items / np.sqrt(np.sum(items**2, axis=1, keepdims=True))
bounds = np.linspace(0, n, n // batch_size + 1, dtype=np.int64)

index = faiss.IndexFlatIP(59)
index.add(items)

distance = np.zeros((n, 4), dtype=np.float32)
index_val = np.zeros((n, 4), dtype=np.int64)
for lower, upper in zip(bounds[:-1], bounds[1:]):
    distance[lower:upper], index_val[lower:upper] = index.search(items[lower:upper], 4)
assert np.sum(index_val[:, 0] != np.arange(n)) == 0
distance = distance[:, 1:]
index_val = index_val[:, 1:]