Einbettungen mit dem integrierten Two-Tower-Algorithmus trainieren

Kunden, die keine eigene Lösung haben, um Einbettungen für Vertex AI Matching Engine zu trainieren, können Einbettungen mit dem integrierten Two-Tower-Algorithmus trainieren.

Überblick

Das Two-Tower-Modell trainiert Einbettungen mithilfe von Daten, die mit Labels versehen wurden. Beim Two-Tower-Modell werden ähnliche Objekttypen wie Nutzerprofile, Suchanfragen, Webdokumente, Antwortdurchläufe oder Bilder im selben Vektorbereich gekoppelt, sodass verwandte Elemente nahe beieinander liegen. Das Two-Tower-Modell besteht aus zwei Encoder-Towern: dem Query-Tower (Abfragen) und dem Candidate-Tower (Kandidaten). Diese Towers enthalten unabhängige Objekte in einem gemeinsamen Einbettungsbereich, in dem Matching Engine ähnlich übereinstimmende Objekte abrufen kann.

Zum Trainieren eines Two-Tower-Modells verwendet Google Paare aus relevanten Elementen. Jedes Paar besteht aus einem Abfragedokument und einem Kandidatendokument. Dokumente enthalten beliebige benutzerdefinierte Features, einschließlich Textfeatures sowie numerische und kategoriale Features. Nach dem Training exportiert der integrierte Two-Tower-Algorithmus zwei TensorFlow-SavedModels – einen Abfrage-Encoder und einen Kandidaten-Encoder, die die entsprechenden Dokumente in Einbettungen konvertieren. Bei einem Abfrageelement verwendet Match Engine den Abfrage-Encoder, um eine Abfrageeinbettung zu erzeugen, und den Index, um ähnliche Kandidateneinbettungen zu finden. Der Kandidaten-Encoder wird von Matching Engine verwendet, um alle Elemente zu indexieren und mit einem Verfahren für die Annäherung an den nächsten Nachbarn zu finden.

Eingabedaten

Der integrierte Two-Tower-Algorithmus erwartet zwei Eingaben:

  • Trainingsdaten: Paare von Dokumenten, die zum Trainieren des Modells verwendet werden. Folgende Dateiformate werden unterstützt:
  • Eingabeschema: Eine JSON-Datei mit dem Schema der Eingabedaten sowie zusätzliche Featurekonfigurationen.

Trainingsdaten

Trainingsdaten bestehen aus Abfragedokumenten und Kandidatendokumenten. Sie müssen nur positive Paare angeben, bei denen die Abfrage- und Kandidatendokumente als Übereinstimmung betrachtet werden.

Abfrage- und Kandidatendokumente bestehen aus mehreren benutzerdefinierten Features. Der Algorithmus unterstützt derzeit die folgenden Featuretypen:

Featuretyp Beschreibung
Text Eine Liste von Strings, die ein Textfeature darstellen.
Id Ein Skalarstring, der die eindeutige ID eines Elements darstellt.
Categorical Eine Liste von Strings, die ein kategoriales Feature darstellen. Die Strings müssen eine Teilmenge einer im Wörterbuch angegebenen Vokabularliste sein.
Numeric Eine Liste der Gleitkommazahlen, die ein numerisches Feature darstellen.
Vector Ein Float-Vektor mit fester Länge, der unverändert eingegeben wird.

JSON Lines-Format

Jede Zeile ist ein JSON-String einer einzelnen Instanz der Trainingsdaten. Das Format ist ein Wörterbuch mit zwei Schlüsseln: query und candidate. Diese wiederum enthalten Schlüssel, die die Features der einzelnen Dokumente darstellen.

Beispiel:

    {
        "query":
        {
            "movie_genre": ["action"],
        },
        "candidate":
        {
            "ratings": [7.3, 6.0, 8.1],
            "movie_name": ["mission impossible"]
        }
    }

TFRecord-Format

Jede Trainingsinstanz ist ein serialisiertes TF-Beispiel. Das Präfix query_ oder candidate_ gibt das Dokument an, zu dem das Feature gehört.

Wir empfehlen die Ausführung der TensorFlow-Anleitung zu TFRecord, um die TFRecord-Datei zu generieren.

Beispiel:

    features {
        feature {
            key: "query_movie_genre"
            value {
                bytes_list {
                    value: "action"
                }
            }
        }
        feature {
            key: "candidate_ratings"
            value {
                float_list {
                    value: [7.3, 6.0, 8.1]
                }
            }
        }
        feature {
            key: "candidate_movie_name"
            value {
                bytes_list {
                    value: "mission impossible"
                }
            }
        }
    }

Eingabeschema

Das Eingabeschema ist eine JSON-Datei, die das Schema der Trainingseingabe zusammen mit den Featurekonfigurationen beschreibt. Wie auch die Trainingsdaten enthält das Eingabeschema zwei Schlüssel: query und candidate, die jeweils die Eingabeschemas der Abfrage und der Kandidatendokumente angeben. Die mit diesen Schlüsseln verknüpften Werte sind ein Wörterbuch, das die Featurenamen ihrem Featuretyp und ihrer Konfiguration zuordnet.

Für die Features Id, Categorical, Text und Vector müssen Sie eine Konfigurationsdatei mit den folgenden zusätzlichen Informationen bereitstellen:

  • Ein Id-Feature erfordert:
    • num_buckets: Anzahl der Hash-Buckets der ID-Funktion. Wenn die Anzahl der IDs klein ist, sollten Sie darauf achten, dass die Anzahl der Buckets der Anzahl der IDs entspricht.
  • Ein Categorical-Feature erfordert:
    • vocab: Eine Liste aller möglichen Werte (Strings) für dieses Feature. Jeder unbekannte kategoriale Wert (außerhalb des Vokabulars), der erkannt wird, wird als separater Wert außerhalb des Vokabulars codiert.
  • Ein Text-Feature kann Folgendes enthalten:
    • embedding_module (optional): Ein String, mit dem das zugrunde liegende Texteinbettungsmodul angegeben wird. Weitere Informationen finden Sie unter Textcodierung.
  • Ein Vector-Feature erfordert:
    • length: Eine positive Ganzzahl, die die Länge (Dimension) der Vektoreingabe darstellt.

Das folgende Beispiel zeigt ein Eingabeschema.

      "query":
       {
          "movie_genre": {
            "feature_type": "Categorical",
            "config": {
              "vocab": ["action", "comedy", "drama"]
            }
          }
       },
       "candidate":
       {
          "ratings": {
            "feature_type": "Numeric"
          },
          "movie_name": {
            "feature_type": "Text",
            "config": {
               "embedding_module":"gs://my_bucket/my_saved_model"
            }
          }
       }

Textcodierung

Textfeatures stellen häufig die wichtigsten Signale für Two-Tower-Modelle bereit. Mit Match Engine können Sie einen Text-Encoder auf mehrere Arten angeben. Sie haben die Möglichkeit, für jedes Textfeature einen anderen Text-Encoder festzulegen. Der Standard-Text-Encoder ist das vortrainierte NNLM-TFHub-Modul.

Textcodierungen akzeptieren als Eingabe einen eindimensionalen Batch von Strings und geben einen zweidimensionalen Float-Tensor mit Einbettungen der Form [batch_size, embedding_dimension] aus. Alle Text-Encoder für den integrierten Two-Tower-Algorithmus, benutzerdefiniert oder nicht, müssen diesem Format folgen. Text-Encodern müssen feste Parameter haben und ihre Gewichtung kann während des Trainings nicht aktualisiert werden.

Für das Texteinbettungsmodul sind folgende drei Optionen verfügbar:

  • TFHub-Pfad zu einem Texteinbettungsmodul, verpackt in einem hub.KerasLayer. Verwenden Sie beispielsweise https://tfhub.dev/google/universal-sentence-encoder/4, um den universal-sentence-encoder anzugeben.
  • Cloud Storage-Pfad zu einem TensorFlow-SavedModel. Dieses SavedModel ist auch in einem hub.KerasLayer verpackt. Es ist wichtig, dass das Projekt Zugriff auf diesen Cloud Storage-Pfad hat.
  • Ein voreingestelltes Einbettungsmodul, das Matching Engine bereitstellt (Name beginnt mit @). Voreinstellungen sind vortrainierte Texteinbettungsmodule, die für ein verteiltes Training im integrierten Two-Tower-Algorithmus verifiziert wurden. Es gibt zwei Voreinstellungen:
    • NNLM (@nnlm, der Standard-Text-Encoder)
    • ALBERT (@albert-base)

Training

Für ein Training mit einem einzelnen Knoten mit dem integrierten Zwei-Tower-Algorithmus verwenden Sie den im Folgenden aufgeführten Befehl. Dieser Befehl erstellt eine CustomJob-Ressource, die eine einzelne CPU-VM-Instanz verwendet. Informationen zu den Flags, die Sie für das Training verwenden können, finden Sie unter Flags. Beim Training mit einem einzelnen Knoten sind die Parameter training_steps_per_epoch und eval_steps_per_epoch optional.

# Set a unique name for the job to run.
JOB_NAME = f'two_tower_cpu_{dataset_name}_{timestamp}'

# URI of the Two-Tower Docker image.
LEARNER_IMAGE_URI = 'us-docker.pkg.dev/vertex-ai-restricted/builtin-algorithm/two-tower'

# The region to run the job in.
REGION = 'us-central1'

# Set the training data and input schema paths.
TRAINING_DATA_PATH = f'gs://cloud-samples-data/vertex-ai/matching-engine/two-tower/{dataset_name}/training_data/*'
INPUT_SCHEMA_PATH = f'gs://cloud-samples-data/vertex-ai/matching-engine/two-tower/{dataset_name}/input_schema.json'

# Set a location for the output.
OUTPUT_DIR = f'gs://{your_bucket_name}/experiment/output'

# Batch size and number of epochs to train.
TRAIN_BATCH_SIZE = 100
NUM_EPOCHS = 10

!gcloud beta ai custom-jobs create \
    --display-name={JOB_NAME} \
    --worker-pool-spec=machine-type=n1-standard-8,replica-count=1,container-image-uri={LEARNER_IMAGE_URI} \
    --region={REGION} \
    --args=--training_data_path={TRAINING_DATA_PATH} \
    --args=--input_schema_path={INPUT_SCHEMA_PATH} \
    --args=--job-dir={OUTPUT_DIR} \
    --args=--train_batch_size={TRAIN_BATCH_SIZE} \
    --args=--num_epochs={NUM_EPOCHS}

Tutorial

Ausführliche Informationen dazu, wie Sie komplexeren Eingaben in den integrierten Two-Tower-Algorithmus einspeisen und mit Vertex AI trainieren, starten Sie das Beispielnotebook in Colab. In diesem Notebook wird auch veranschaulicht, wie Sie auf GPUs trainieren und Details und Logs in TensorBoard anzeigen lassen, während Ihr Job ausgeführt wird.

Flags

Verwenden Sie beim Trainieren eines Modells mit dem integrierten Two-Tower-Modell die folgenden allgemeinen und modellspezifischen Two-Tower-Trainings-Flags.

Allgemeine Trainings-Flags

Die im Folgenden aufgeführten benutzerdefinierten Trainings-Flags werden am häufigsten verwendet. Weitere Informationen finden Sie unter Benutzerdefinierte Trainingsjobs erstellen.

  • worker-pool-spec: Die vom benutzerdefinierten Job verwendete Konfiguration des Worker-Pools. Wenn Sie einen benutzerdefinierten Job mit mehreren Worker-Pools erstellen möchten, geben Sie mehrere worker-pool-spec-Konfigurationen an.

    Eine worker-pool-spec kann die im Folgenden aufgeführten Felder enthalten, die mit entsprechenden Feldern in der WorkerPoolSpec API-Nachricht aufgeführt sind.

    • machine-type: Der Maschinentyp für den Pool. Eine Liste der unterstützten Maschinen finden Sie unter Maschinentypen.
    • replica-count: Die Anzahl der Replikate der Maschine im Pool.
    • container-image-uri: Das Docker-Image, das auf jedem Worker ausgeführt werden soll. Wenn Sie den integrierten Two-Tower-Algorithmus verwenden möchten, muss für das Docker-Image us-docker.pkg.dev/vertex-ai-restricted/builtin-algorithm/two-tower:latest festgelegt sein.
  • display-name: Der Name des Jobs.

  • region: Die Region, in der der Job ausgeführt werden soll.

Modellspezifische Trainings-Flags für Two-Tower

  • training_data_path: Das Muster für den Cloud Storage-URI, der angibt, wo Trainingsdaten gespeichert werden (z. B. gs://<bucket>/<folder>/*).
  • eval_data_path: Das Muster für den Cloud Storage-URI, der angibt, wo Auswertungsdaten gespeichert werden (z. B. gs://<bucket>/<folder>/*).
  • input_schema_path: Der Cloud Storage-URI, der angibt, wo das JSON-Eingabeschema gespeichert ist (z. B. gs://<bucket>/<folder>/<schema_file>)
  • input_file_format: Das Format der Eingabedatei. input_file_format kann eines der folgenden Formate haben:
    • jsonl (Standard)
    • tfrecord
  • job-dir: Der Cloud Storage-URI eines Verzeichnisses, in dem Sie Trainingsausgaben speichern möchten.
  • candidate_data_path: Das Muster für den Cloud Storage-URI, der angibt, wo Kandidatendaten gespeichert werden (z. B. gs://<bucket>/<folder>/*). candidate_data_path ist nur erforderlich, wenn Nicht-Standard-top_k_categorical_accuracy-Messwerte anstelle von auc verwendet werden. Wenn candidate_data_path nicht festgelegt ist, wird der Kandidatensatz aus den Trainings- und Auswertungsdaten erstellt.
  • train_batch_size: Die Batchgröße für das Training. Der Standardwert ist 100.
  • eval_batch_size: Die Batchgröße für die Bewertung. Der Standardwert ist 100.
  • eval_split: Der für das Auswertungs-Dataset zu verwendende Split, wenn eval_data_path nicht angegeben ist. Der Standardwert ist 0.2 (20 % der Daten aus training_data_path werden als Auswertungsdaten verwendet).
  • metrics: Die Messwerte, die zum Bewerten des Modells verwendet werden. Sie können Evaluierungsmesswerte aus den Logdateien abrufen oder in TensorBoard hochladen, insbesondere aus gs://<job-dir>/tensorboard/validation, wie in der TensorBoard-Dokumentation beschrieben. metrics kann eines der Folgenden sein:

    • auc (Standard)
    • top_k_categorical_accuracy
    • precision_at_1
  • optimizer: Das zum Trainieren des Modells verwendete Optimierungstool. Sie können den Namen eines beliebigen TensorFlow 2.3 Keras-Optimierungstools in Kleinbuchstaben verwenden, z.B. sgd, adam oder ftrl. Weitere Informationen finden Sie unter Modul: tf.keras.Optimizers. Der Standardwert ist adagrad.

  • learning_rate: Die ML-Rate (maschinelles Lernen) für das Training mit dem Optimierungstool.

  • momentum: Der Impuls für das Optimierungstool für sgd oder rmsprop, sofern angegeben.

  • num_epochs: Die Anzahl der Epochen für das Training. Der Standardwert ist 10.

  • num_hidden_layers: Die Anzahl der ausgeblendeten Ebenen. Standardmäßig sind keine ausgeblendeten Ebenen vorhanden.

  • num_nodes_hidden_layer{index}: Die Anzahl der Knoten für die ausgeblendete Ebene mit diesem bestimmten Index. Verwenden Sie beispielsweise für die erste verborgene Ebene num_nodes_hidden_layer1. Der Indexbereich reicht von 1 bis 20.

  • output_dim: Eine einzelne Ganzzahl, die die Ausgabecodierungsdimension für jeden Turm des Zwei-Turm-Modells darstellt. Der Standardwert ist 64.

  • cache_input: Ein boolescher Wert, mit dem angegeben wird, ob Trainings- und Auswertungsdaten im Cache gespeichert werden sollen. Dieses Flag wird nur für kleine Datasets empfohlen.

  • training_steps_per_epoch: Die Anzahl der Schritte pro Epoche, für die das Training ausgeführt werden soll. Dieses Flag wird nur benötigt, wenn Sie mehrere Knoten für das Training verwenden oder wenn Ihr primärer Trainingsknoten mehr als 1 GPU hat. Der Standardwert ist None.

  • eval_steps_per_epoch: Die Anzahl der Schritte pro Epoche, für die die Bewertung ausgeführt werden soll. Dieses Flag wird nur benötigt, wenn Sie mehrere Knoten für das Training verwenden oder Ihr primärer Trainingsknoten mehr als 1 GPU hat. Der Standardwert ist None.

  • gpu_memory_alloc: Die Menge des Arbeitsspeichers in Megabyte, die pro GPU zugewiesen wird. Standardmäßig gilt kein Limit.

Nächste Schritte