Посмотреть на TensorFlow.org | Запускаем в Google Colab | Посмотреть исходный код на GitHub | Скачать блокнот |
Retrieval модели часто строятся на поверхности горсть лучших кандидатов из миллионов или даже сотни миллионов кандидатов. Чтобы иметь возможность реагировать на контекст и поведение пользователя, они должны уметь делать это «на лету» за считанные миллисекунды.
Приблизительный поиск ближайшего соседа (ANN) - это технология, которая делает это возможным. В этом руководстве мы покажем, как использовать ScaNN - современный пакет поиска ближайшего соседа - для плавного масштабирования извлечения TFRS до миллионов элементов.
Что такое ScaNN?
ScaNN - это библиотека от Google Research, которая выполняет крупномасштабный поиск схожести векторов. Имея базу данных возможных встраиваний, ScaNN индексирует эти вложения таким образом, чтобы их можно было быстро найти во время вывода. ScaNN использует современные методы сжатия векторов и тщательно реализованные алгоритмы для достижения наилучшего компромисса между скоростью и точностью. Он может значительно превзойти поиск методом грубой силы, при этом мало жертвуя с точки зрения точности.
Построение модели на основе ScaNN
Для того, чтобы попробовать ScaNN в ОКФЕ, мы создадим простую MovieLens поисковую модель, так же , как мы это делали в основных поисковых обучающей программе . Если вы следовали этому руководству, этот раздел будет вам знаком, и его можно смело пропустить.
Для начала установите TFRS и TensorFlow Datasets:
pip install -q tensorflow-recommenders
pip install -q --upgrade tensorflow-datasets
Нам также необходимо установить scann
: это необязательная зависимость ОКФ, и поэтому должен быть установлен отдельно.
pip install -q scann
Настроить весь необходимый импорт.
from typing import Dict, Text
import os
import pprint
import tempfile
import numpy as np
import tensorflow as tf
import tensorflow_datasets as tfds
import tensorflow_recommenders as tfrs
И загрузите данные:
# Load the MovieLens 100K data.
ratings = tfds.load(
"movielens/100k-ratings",
split="train"
)
# Get the ratings data.
ratings = (ratings
# Retain only the fields we need.
.map(lambda x: {"user_id": x["user_id"], "movie_title": x["movie_title"]})
# Cache for efficiency.
.cache(tempfile.NamedTemporaryFile().name)
)
# Get the movies data.
movies = tfds.load("movielens/100k-movies", split="train")
movies = (movies
# Retain only the fields we need.
.map(lambda x: x["movie_title"])
# Cache for efficiency.
.cache(tempfile.NamedTemporaryFile().name))
2021-10-02 11:53:59.413405: E tensorflow/stream_executor/cuda/cuda_driver.cc:271] failed call to cuInit: CUDA_ERROR_NO_DEVICE: no CUDA-capable device is detected
Прежде чем мы сможем построить модель, нам нужно настроить словари пользователей и фильмов:
user_ids = ratings.map(lambda x: x["user_id"])
unique_movie_titles = np.unique(np.concatenate(list(movies.batch(1000))))
unique_user_ids = np.unique(np.concatenate(list(user_ids.batch(1000))))
2021-10-02 11:54:00.296290: W tensorflow/core/kernels/data/cache_dataset_ops.cc:233] The calling iterator did not fully read the dataset being cached. In order to avoid unexpected truncation of the dataset, the partially cached contents of the dataset will be discarded. This can happen if you have an input pipeline similar to `dataset.cache().take(k).repeat()`. You should use `dataset.take(k).cache().repeat()` instead. 2021-10-02 11:54:04.003150: W tensorflow/core/kernels/data/cache_dataset_ops.cc:233] The calling iterator did not fully read the dataset being cached. In order to avoid unexpected truncation of the dataset, the partially cached contents of the dataset will be discarded. This can happen if you have an input pipeline similar to `dataset.cache().take(k).repeat()`. You should use `dataset.take(k).cache().repeat()` instead.
Также настроим обучающий и тестовый наборы:
tf.random.set_seed(42)
shuffled = ratings.shuffle(100_000, seed=42, reshuffle_each_iteration=False)
train = shuffled.take(80_000)
test = shuffled.skip(80_000).take(20_000)
Определение модели
Так же , как и в основных поисковых обучающей программе , мы создадим простую модель две башни.
class MovielensModel(tfrs.Model):
def __init__(self):
super().__init__()
embedding_dimension = 32
# Set up a model for representing movies.
self.movie_model = tf.keras.Sequential([
tf.keras.layers.StringLookup(
vocabulary=unique_movie_titles, mask_token=None),
# We add an additional embedding to account for unknown tokens.
tf.keras.layers.Embedding(len(unique_movie_titles) + 1, embedding_dimension)
])
# Set up a model for representing users.
self.user_model = tf.keras.Sequential([
tf.keras.layers.StringLookup(
vocabulary=unique_user_ids, mask_token=None),
# We add an additional embedding to account for unknown tokens.
tf.keras.layers.Embedding(len(unique_user_ids) + 1, embedding_dimension)
])
# Set up a task to optimize the model and compute metrics.
self.task = tfrs.tasks.Retrieval(
metrics=tfrs.metrics.FactorizedTopK(
candidates=movies.batch(128).cache().map(self.movie_model)
)
)
def compute_loss(self, features: Dict[Text, tf.Tensor], training=False) -> tf.Tensor:
# We pick out the user features and pass them into the user model.
user_embeddings = self.user_model(features["user_id"])
# And pick out the movie features and pass them into the movie model,
# getting embeddings back.
positive_movie_embeddings = self.movie_model(features["movie_title"])
# The task computes the loss and the metrics.
return self.task(user_embeddings, positive_movie_embeddings, compute_metrics=not training)
Примерка и оценка
Модель TFRS - это просто модель Кераса. Мы можем его скомпилировать:
model = MovielensModel()
model.compile(optimizer=tf.keras.optimizers.Adagrad(learning_rate=0.1))
Оцените это:
model.fit(train.batch(8192), epochs=3)
Epoch 1/3 10/10 [==============================] - 3s 223ms/step - factorized_top_k/top_1_categorical_accuracy: 0.0000e+00 - factorized_top_k/top_5_categorical_accuracy: 0.0000e+00 - factorized_top_k/top_10_categorical_accuracy: 0.0000e+00 - factorized_top_k/top_50_categorical_accuracy: 0.0000e+00 - factorized_top_k/top_100_categorical_accuracy: 0.0000e+00 - loss: 69808.9716 - regularization_loss: 0.0000e+00 - total_loss: 69808.9716 Epoch 2/3 10/10 [==============================] - 3s 222ms/step - factorized_top_k/top_1_categorical_accuracy: 0.0000e+00 - factorized_top_k/top_5_categorical_accuracy: 0.0000e+00 - factorized_top_k/top_10_categorical_accuracy: 0.0000e+00 - factorized_top_k/top_50_categorical_accuracy: 0.0000e+00 - factorized_top_k/top_100_categorical_accuracy: 0.0000e+00 - loss: 67485.8842 - regularization_loss: 0.0000e+00 - total_loss: 67485.8842 Epoch 3/3 10/10 [==============================] - 3s 220ms/step - factorized_top_k/top_1_categorical_accuracy: 0.0000e+00 - factorized_top_k/top_5_categorical_accuracy: 0.0000e+00 - factorized_top_k/top_10_categorical_accuracy: 0.0000e+00 - factorized_top_k/top_50_categorical_accuracy: 0.0000e+00 - factorized_top_k/top_100_categorical_accuracy: 0.0000e+00 - loss: 66311.9581 - regularization_loss: 0.0000e+00 - total_loss: 66311.9581 <keras.callbacks.History at 0x7fc02423c150>
И оцените это.
model.evaluate(test.batch(8192), return_dict=True)
3/3 [==============================] - 2s 246ms/step - factorized_top_k/top_1_categorical_accuracy: 0.0011 - factorized_top_k/top_5_categorical_accuracy: 0.0095 - factorized_top_k/top_10_categorical_accuracy: 0.0222 - factorized_top_k/top_50_categorical_accuracy: 0.1261 - factorized_top_k/top_100_categorical_accuracy: 0.2363 - loss: 49466.8789 - regularization_loss: 0.0000e+00 - total_loss: 49466.8789 {'factorized_top_k/top_1_categorical_accuracy': 0.0010999999940395355, 'factorized_top_k/top_5_categorical_accuracy': 0.009549999609589577, 'factorized_top_k/top_10_categorical_accuracy': 0.022199999541044235, 'factorized_top_k/top_50_categorical_accuracy': 0.1261499971151352, 'factorized_top_k/top_100_categorical_accuracy': 0.23634999990463257, 'loss': 28242.8359375, 'regularization_loss': 0, 'total_loss': 28242.8359375}
Примерный прогноз
Самый простой способ получить лучших кандидатов в ответ на запрос - это сделать это с помощью грубой силы: вычислить оценки пользовательских фильмов для всех возможных фильмов, отсортировать их и выбрать пару лучших рекомендаций.
В СКР, это достигается с помощью BruteForce
слоя:
brute_force = tfrs.layers.factorized_top_k.BruteForce(model.user_model)
brute_force.index_from_dataset(
movies.batch(128).map(lambda title: (title, model.movie_model(title)))
)
<tensorflow_recommenders.layers.factorized_top_k.BruteForce at 0x7fbfc1d4fe10>
После того, как создаются и заполняются с кандидатами ( с помощью index
метода), мы можем назвать его , чтобы получить предсказания из:
# Get predictions for user 42.
_, titles = brute_force(np.array(["42"]), k=3)
print(f"Top recommendations: {titles[0]}")
Top recommendations: [b'Homeward Bound: The Incredible Journey (1993)' b"Kid in King Arthur's Court, A (1995)" b'Rudy (1993)']
Для небольшого набора данных из менее 1000 фильмов это очень быстро:
%timeit _, titles = brute_force(np.array(["42"]), k=3)
983 µs ± 5.44 µs per loop (mean ± std. dev. of 7 runs, 1000 loops each)
Но что будет, если у нас будет больше кандидатов - миллионы вместо тысяч?
Мы можем смоделировать это, индексируя все наши фильмы несколько раз:
# Construct a dataset of movies that's 1,000 times larger. We
# do this by adding several million dummy movie titles to the dataset.
lots_of_movies = tf.data.Dataset.concatenate(
movies.batch(4096),
movies.batch(4096).repeat(1_000).map(lambda x: tf.zeros_like(x))
)
# We also add lots of dummy embeddings by randomly perturbing
# the estimated embeddings for real movies.
lots_of_movies_embeddings = tf.data.Dataset.concatenate(
movies.batch(4096).map(model.movie_model),
movies.batch(4096).repeat(1_000)
.map(lambda x: model.movie_model(x))
.map(lambda x: x * tf.random.uniform(tf.shape(x)))
)
Мы можем построить BruteForce
индекс этого большего набора данных:
brute_force_lots = tfrs.layers.factorized_top_k.BruteForce()
brute_force_lots.index_from_dataset(
tf.data.Dataset.zip((lots_of_movies, lots_of_movies_embeddings))
)
<tensorflow_recommenders.layers.factorized_top_k.BruteForce at 0x7fbfc1d80610>
Рекомендации все те же
_, titles = brute_force_lots(model.user_model(np.array(["42"])), k=3)
print(f"Top recommendations: {titles[0]}")
Top recommendations: [b'Homeward Bound: The Incredible Journey (1993)' b"Kid in King Arthur's Court, A (1995)" b'Rudy (1993)']
Но на это уходит гораздо больше времени. С набором кандидатов в 1 миллион фильмов предсказание методом грубой силы становится довольно медленным:
%timeit _, titles = brute_force_lots(model.user_model(np.array(["42"])), k=3)
33 ms ± 245 µs per loop (mean ± std. dev. of 7 runs, 10 loops each)
По мере роста числа кандидатов количество необходимого времени растет линейно: с 10 миллионами кандидатов обслуживание лучших кандидатов займет 250 миллисекунд. Это явно слишком медленно для живого сервиса.
Вот тут и вступают в игру приблизительные механизмы.
Использование ScaNN в СКР осуществляется через tfrs.layers.factorized_top_k.ScaNN
слоя. Он следует тому же интерфейсу, что и другие верхние k слоев:
scann = tfrs.layers.factorized_top_k.ScaNN(num_reordering_candidates=100)
scann.index_from_dataset(
tf.data.Dataset.zip((lots_of_movies, lots_of_movies_embeddings))
)
<tensorflow_recommenders.layers.factorized_top_k.ScaNN at 0x7fbfc2571990>
Рекомендации (примерно!) Одинаковые
_, titles = scann(model.user_model(np.array(["42"])), k=3)
print(f"Top recommendations: {titles[0]}")
Top recommendations: [b'Homeward Bound: The Incredible Journey (1993)' b"Kid in King Arthur's Court, A (1995)" b'Rudy (1993)']
Но их вычислить намного быстрее:
%timeit _, titles = scann(model.user_model(np.array(["42"])), k=3)
4.35 ms ± 34.7 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)
В этом случае мы можем получить 3 лучших фильма из набора ~ 1 миллиона примерно за 2 миллисекунды: в 15 раз быстрее, чем при вычислении лучших кандидатов с помощью грубой силы. Преимущество приближенных методов становится еще больше для больших наборов данных.
Оценка приближения
При использовании приблизительных механизмов поиска K (таких как ScaNN) скорость поиска часто достигается за счет точности. Чтобы понять этот компромисс, важно измерить метрики оценки модели при использовании ScaNN и сравнить их с базовыми показателями.
К счастью, TFRS упрощает эту задачу. Мы просто переопределяем метрики в задаче извлечения метриками с использованием ScaNN, повторно компилируем модель и запускаем оценку.
Чтобы провести сравнение, давайте сначала запустим базовые результаты. Нам все еще нужно переопределить наши метрики, чтобы убедиться, что они используют расширенный набор кандидатов, а не исходный набор фильмов:
# Override the existing streaming candidate source.
model.task.factorized_metrics = tfrs.metrics.FactorizedTopK(
candidates=lots_of_movies_embeddings
)
# Need to recompile the model for the changes to take effect.
model.compile()
%time baseline_result = model.evaluate(test.batch(8192), return_dict=True, verbose=False)
CPU times: user 22min 5s, sys: 2min 7s, total: 24min 12s Wall time: 51.9 s
Мы можем сделать то же самое с помощью ScaNN:
model.task.factorized_metrics = tfrs.metrics.FactorizedTopK(
candidates=scann
)
model.compile()
# We can use a much bigger batch size here because ScaNN evaluation
# is more memory efficient.
%time scann_result = model.evaluate(test.batch(8192), return_dict=True, verbose=False)
CPU times: user 10.5 s, sys: 3.26 s, total: 13.7 s Wall time: 1.85 s
Оценка на основе ScaNN намного быстрее: она в десять раз быстрее! Это преимущество станет еще больше для больших наборов данных, поэтому для больших наборов данных может быть разумным всегда запускать оценку на основе ScaNN, чтобы повысить скорость разработки модели.
Но как насчет результатов? К счастью, в этом случае результаты почти такие же:
print(f"Brute force top-100 accuracy: {baseline_result['factorized_top_k/top_100_categorical_accuracy']:.2f}")
print(f"ScaNN top-100 accuracy: {scann_result['factorized_top_k/top_100_categorical_accuracy']:.2f}")
Brute force top-100 accuracy: 0.15 ScaNN top-100 accuracy: 0.27
Это говорит о том, что в этом искусственном наборе данных есть небольшая потеря от аппроксимации. В общем, все приближенные методы демонстрируют компромисс между скоростью и точностью. Чтобы понять это более подробно вы можете проверить Эрика Bernhardsson в ИНС тестов .
Развертывание примерной модели
ScaNN
основанное модель полностью интегрирована в модели TensorFlow, и служить так же легко , как служить любой другой модели TensorFlow.
Мы можем сохранить его как SavedModel
объект
lots_of_movies_embeddings
<ConcatenateDataset shapes: (None, 32), types: tf.float32>
# We re-index the ScaNN layer to include the user embeddings in the same model.
# This way we can give the saved model raw features and get valid predictions
# back.
scann = tfrs.layers.factorized_top_k.ScaNN(model.user_model, num_reordering_candidates=1000)
scann.index_from_dataset(
tf.data.Dataset.zip((lots_of_movies, lots_of_movies_embeddings))
)
# Need to call it to set the shapes.
_ = scann(np.array(["42"]))
with tempfile.TemporaryDirectory() as tmp:
path = os.path.join(tmp, "model")
tf.saved_model.save(
scann,
path,
options=tf.saved_model.SaveOptions(namespace_whitelist=["Scann"])
)
loaded = tf.saved_model.load(path)
2021-10-02 11:55:53.875291: W tensorflow/python/util/util.cc:348] Sets are not currently considered sequences, but this may change in the future, so consider avoiding using them. WARNING:absl:Found untraced functions such as query_with_exclusions while saving (showing 1 of 1). These functions will not be directly callable after loading. INFO:tensorflow:Assets written to: /tmp/tmpm0piq8hx/model/assets INFO:tensorflow:Assets written to: /tmp/tmpm0piq8hx/model/assets
а затем загрузите его и подавайте, возвращая точно такие же результаты:
_, titles = loaded(tf.constant(["42"]))
print(f"Top recommendations: {titles[0][:3]}")
Top recommendations: [b'Homeward Bound: The Incredible Journey (1993)' b"Kid in King Arthur's Court, A (1995)" b'Rudy (1993)']
Полученную модель можно использовать в любом сервисе Python, на котором установлены TensorFlow и ScaNN.
Он также может быть подан с использованием специальной версии TensorFlow сервировки, доступной в качестве контейнера Докера на Докер Hub . Вы также можете построить сами изображения из Dockerfile .
Тюнинг ScaNN
Теперь давайте рассмотрим настройку нашего слоя ScaNN, чтобы получить лучший компромисс между производительностью и точностью. Чтобы сделать это эффективно, нам сначала нужно измерить нашу базовую производительность и точность.
Сверху, у нас уже есть измерение задержки нашей модели для обработки одного (не пакетного) запроса (хотя обратите внимание, что значительная часть этой задержки связана с компонентами модели, не относящимися к ScaNN).
Теперь нам нужно исследовать точность ScaNN, которую мы измеряем с помощью отзыва. Отзыв @ k, равный x%, означает, что если мы используем грубую силу для получения истинных k верхних соседей и сравниваем эти результаты с использованием ScaNN для получения также верхних k соседей, x% результатов ScaNN будет в истинных результатах грубой силы. Давайте посчитаем отзыв для текущего искателя ScaNN.
Во-первых, нам нужно сгенерировать грубую силу, наземную истину top-k:
# Process queries in groups of 1000; processing them all at once with brute force
# may lead to out-of-memory errors, because processing a batch of q queries against
# a size-n dataset takes O(nq) space with brute force.
titles_ground_truth = tf.concat([
brute_force_lots(queries, k=10)[1] for queries in
test.batch(1000).map(lambda x: model.user_model(x["user_id"]))
], axis=0)
Наша переменная titles_ground_truth
теперь содержит рекомендации фильмов топ-10 возвращаемых поиск перебора. Теперь мы можем вычислить те же рекомендации при использовании ScaNN:
# Get all user_id's as a 1d tensor of strings
test_flat = np.concatenate(list(test.map(lambda x: x["user_id"]).batch(1000).as_numpy_iterator()), axis=0)
# ScaNN is much more memory efficient and has no problem processing the whole
# batch of 20000 queries at once.
_, titles = scann(test_flat, k=10)
Затем мы определяем нашу функцию, которая вычисляет отзыв. Для каждого запроса он подсчитывает, сколько результатов находится на пересечении результатов грубой силы и результатов ScaNN, и делит это на количество результатов грубой силы. Среднее значение этого количества по всем запросам - это наш отзыв.
def compute_recall(ground_truth, approx_results):
return np.mean([
len(np.intersect1d(truth, approx)) / len(truth)
for truth, approx in zip(ground_truth, approx_results)
])
Это дает нам базовый отзыв @ 10 с текущей конфигурацией ScaNN:
print(f"Recall: {compute_recall(titles_ground_truth, titles):.3f}")
Recall: 0.931
Мы также можем измерить базовую задержку:
%timeit -n 1000 scann(np.array(["42"]), k=10)
4.67 ms ± 25 µs per loop (mean ± std. dev. of 7 runs, 1000 loops each)
Посмотрим, сможем ли мы сделать лучше!
Для этого нам нужна модель того, как ручки настройки ScaNN влияют на производительность. Наша текущая модель использует алгоритм ScaNN tree-AH. Этот алгоритм разбивает базу данных вложений («дерево»), а затем оценивает наиболее многообещающее из этих разделов с помощью AH, который представляет собой высокооптимизированную процедуру вычисления приблизительного расстояния.
Параметры по умолчанию для TensorFlow рекомендателей ScaNN Keras наборы слоев num_leaves=100
и num_leaves_to_search=10
. Это означает, что наша база данных разделена на 100 непересекающихся подмножеств, и 10 наиболее многообещающих из этих разделов оцениваются с помощью AH. Это означает, что 10/100 = 10% набора данных ищется с помощью AH.
Если мы, скажем, num_leaves=1000
и num_leaves_to_search=100
, мы бы также искать 10% базы данных с АГ. Однако, по сравнению с предыдущей установкой 10% мы будем искать кандидат будет содержать более высокое качество, потому что более высокий num_leaves
позволяет сделать более мелкозернистое решение о том, каких частях набора данных стоит поиск.
Это не удивительно то , что с num_leaves=1000
и num_leaves_to_search=100
мы получаем значительно более высокий отзыв:
scann2 = tfrs.layers.factorized_top_k.ScaNN(
model.user_model,
num_leaves=1000,
num_leaves_to_search=100,
num_reordering_candidates=1000)
scann2.index_from_dataset(
tf.data.Dataset.zip((lots_of_movies, lots_of_movies_embeddings))
)
_, titles2 = scann2(test_flat, k=10)
print(f"Recall: {compute_recall(titles_ground_truth, titles2):.3f}")
Recall: 0.966
Однако, как компромисс, наша задержка также увеличилась. Это потому, что этап разделения стал более дорогим; scann
выбирает верхней 10 из 100 разделов , в то время как scann2
берет верх 100 1000 разделов. Последнее может быть дороже, потому что требует увеличения количества разделов в 10 раз.
%timeit -n 1000 scann2(np.array(["42"]), k=10)
4.86 ms ± 21.5 µs per loop (mean ± std. dev. of 7 runs, 1000 loops each)
В общем, настройка поиска ScaNN заключается в выборе правильных компромиссов. Изменение каждого отдельного параметра обычно не делает поиск более быстрым и точным; наша цель - настроить параметры, чтобы добиться оптимального баланса между этими двумя конфликтующими целями.
В нашем случае, scann2
значительно улучшилась по сравнению вспомнить scann
определенной цены в латентности. Можем ли мы повернуть назад некоторые другие регуляторы, чтобы сократить задержку, сохранив большую часть нашего преимущества отзыва?
Давайте попробуем найти 70/1000 = 7% набора данных с помощью AH и восстановить только последние 400 кандидатов:
scann3 = tfrs.layers.factorized_top_k.ScaNN(
model.user_model,
num_leaves=1000,
num_leaves_to_search=70,
num_reordering_candidates=400)
scann3.index_from_dataset(
tf.data.Dataset.zip((lots_of_movies, lots_of_movies_embeddings))
)
_, titles3 = scann3(test_flat, k=10)
print(f"Recall: {compute_recall(titles_ground_truth, titles3):.3f}")
Recall: 0.957
scann3
поставляет об абсолютной выгоде вспомнить 3% по сравнению с scann
, а также поставляя более низкое время ожидания:
%timeit -n 1000 scann3(np.array(["42"]), k=10)
4.58 ms ± 37.1 µs per loop (mean ± std. dev. of 7 runs, 1000 loops each)
Эти ручки можно дополнительно отрегулировать для оптимизации для различных точек на границе Парето с точки зрения точности и производительности. Алгоритмы ScaNN позволяют достичь высочайшего уровня производительности в широком диапазоне задач отзыва.
дальнейшее чтение
ScaNN использует передовые методы векторного квантования и высокооптимизированную реализацию для достижения своих результатов. Область векторного квантования имеет богатую историю с множеством подходов. Текущий метод квантования ScaNN подробно рассказывается в данной статье , опубликованной в ICML 2020. Этот документ был также выпущен вместе с этой статьей блог , который дает высокоуровневый обзор нашей техники.
Многие методы связанных квантований упоминаются в ссылках нашей газеты ICML 2020, и другие исследования , связанные с ScaNN перечислены в http://sanjivk.com/