 TensorFlow.orgで表示 |  GoogleColabで実行 |  GitHubでソースを表示 |  ノートブックをダウンロード |
検索モデルは、多くの場合、数百万人かの候補者の何百万人もの数百のうち、上位候補の一握りを表面に構築されています。ユーザーのコンテキストと動作に対応できるようにするには、ユーザーは数ミリ秒でこれをオンザフライで実行できる必要があります。
近似最近傍探索（ANN）は、これを可能にするテクノロジーです。このチュートリアルでは、ScaNN（最先端の最近傍検索パッケージ）を使用して、TFRS検索を数百万のアイテムにシームレスにスケーリングする方法を示します。
ScaNNとは何ですか？
ScaNNは、大規模な高密度ベクトル類似性検索を実行するGoogleResearchのライブラリです。候補埋め込みのデータベースが与えられると、ScaNNは、推論時に迅速に検索できるように、これらの埋め込みにインデックスを付けます。 ScaNNは、最先端のベクトル圧縮技術と慎重に実装されたアルゴリズムを使用して、最高の速度と精度のトレードオフを実現します。精度の点でほとんど犠牲にせずに、ブルートフォース検索を大幅に上回ることができます。
ScaNNを利用したモデルの構築
TFRSでScaNNを試すために、私たちは私たちがやったように、簡単なMovieLens検索モデルを構築します基本的な検索のチュートリアル。そのチュートリアルに従っている場合、このセクションはおなじみであり、安全にスキップできます。
開始するには、TFRSおよびTensorFlowデータセットをインストールします。
pip install -q tensorflow-recommenderspip install -q --upgrade tensorflow-datasets
また、インストールする必要がありscann ：それはTFRSのオプションの依存だし、そう別途インストールする必要があります。
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モデルは単なるKerasモデルです。私たちはそれをコンパイルすることができます：
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}
おおよその予測
クエリに応答して上位の候補を取得する最も簡単な方法は、ブルートフォースを介してそれを行うことです。すべての可能な映画のユーザー映画スコアを計算し、それらを並べ替えて、いくつかの上位の推奨事項を選択します。
TFRSでは、これは、を介して達成される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)']
しかし、彼らははるかに時間がかかります。 100万本の映画の候補セットがあると、ブルートフォース予測は非常に遅くなります。
%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)
候補者の数が増えると、必要な時間は直線的に増加します。1,000万人の候補者がいる場合、上位の候補者にサービスを提供するには250ミリ秒かかります。これは明らかにライブサービスには遅すぎます。
これがおおよそのメカニズムの出番です。
TFRSで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)
この場合、約2ミリ秒で約100万本のセットから上位3本の映画を取得できます。これは、ブルートフォースで最適な候補を計算するよりも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ベースの評価は、はるかに高速です。10倍以上高速です。この利点は、データセットが大きいほどさらに大きくなるため、データセットが大きい場合は、モデル開発の速度を向上させるために、常に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のチェックアウトすることができANNベンチマークを。
近似モデルの展開
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)']
結果のモデルは、TensorFlowとScaNNがインストールされている任意のPythonサービスで提供できます。
また、上のドッカーコンテナとして利用できるTensorFlowのサービングのカスタマイズバージョン、使用して提供することができますドッカーハブ。また、から画像を自分で構築することができDockerfile 。
ScaNNのチューニング
次に、ScaNNレイヤーを調整して、パフォーマンスと精度のトレードオフを改善する方法を見てみましょう。これを効果的に行うには、最初にベースラインのパフォーマンスと精度を測定する必要があります。
上記から、単一の（バッチ処理されていない）クエリを処理するためのモデルのレイテンシーの測定値がすでにあります（ただし、このレイテンシーのかなりの量はモデルの非ScaNNコンポーネントからのものであることに注意してください）。
次に、リコールによって測定するScaNNの精度を調査する必要があります。 x％のrecall @ kは、ブルートフォースを使用して真の上位k近傍を取得し、それらの結果をScaNNを使用して上位k近傍も取得する場合と比較すると、ScaNNの結果のx％が真のブルートフォース結果に含まれることを意味します。現在の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)
])
これにより、現在のScaNN構成でベースラインrecall @ 10が得られます。
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のツリー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 、我々はまた、AHとデータベースの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しながら、100のパーティションのトップ10を選ぶscann2 1000のパーティションのトップ100を選びます。後者は、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検索の調整は、適切なトレードオフを選択することです。個々のパラメータを変更しても、通常、検索がより速く、より正確になるわけではありません。私たちの目標は、これら2つの相反する目標の間で最適にトレードオフするようにパラメーターを調整することです。
我々の場合には、 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/