Entraîner un modèle TensorFlow avec Keras sur Google Kubernetes Engine

La section suivante fournit un exemple d'ajustement d'un modèle BERT pour la classification de séquences à l'aide de la bibliothèque Hugging Face Transformers avec TensorFlow. L'ensemble de données est téléchargé dans un volume Parallelstore associé, ce qui permet à l'entraînement du modèle de lire directement les données du volume.

Prérequis

• Assurez-vous que votre nœud dispose d'au moins 8 Go de mémoire disponible.
• Créez un objet PersistentVolumeClaim demandant un volume basé sur Parallelstore.

Enregistrez le fichier manifeste YAML (parallelstore-csi-job-example.yaml) suivant pour votre tâche d'entraînement de modèle.

  apiVersion: batch/v1
  kind: Job
  metadata:
    name: parallelstore-csi-job-example
  spec:
    template:
      metadata:
        annotations:
            gke-parallelstore/cpu-limit: "0"
            gke-parallelstore/memory-limit: "0"
      spec:
        securityContext:
          runAsUser: 1000
          runAsGroup: 100
          fsGroup: 100
        containers:
        - name: tensorflow
          image: jupyter/tensorflow-notebook@sha256:173f124f638efe870bb2b535e01a76a80a95217e66ed00751058c51c09d6d85d
          command: ["bash", "-c"]
          args:
          - |
            pip install transformers datasets
            python - <<EOF
            from datasets import load_dataset
            dataset = load_dataset("glue", "cola", cache_dir='/data')
            dataset = dataset["train"]
            from transformers import AutoTokenizer
            import numpy as np
            tokenizer = AutoTokenizer.from_pretrained("bert-base-cased")
            tokenized_data = tokenizer(dataset["sentence"], return_tensors="np", padding=True)
            tokenized_data = dict(tokenized_data)
            labels = np.array(dataset["label"])
            from transformers import TFAutoModelForSequenceClassification
            from tensorflow.keras.optimizers import Adam
            model = TFAutoModelForSequenceClassification.from_pretrained("bert-base-cased")
            model.compile(optimizer=Adam(3e-5))
            model.fit(tokenized_data, labels)
            EOF
          volumeMounts:
          - name: parallelstore-volume
            mountPath: /data
        volumes:
        - name: parallelstore-volume
          persistentVolumeClaim:
            claimName: parallelstore-pvc
        restartPolicy: Never
    backoffLimit: 1

Appliquez le fichier manifeste YAML au cluster.

kubectl apply -f parallelstore-csi-job-example.yaml

Vérifiez la progression du chargement des données et de l'entraînement du modèle à l'aide de la commande suivante:

POD_NAME=$(kubectl get pod | grep 'parallelstore-csi-job-example' | awk '{print $1}')
kubectl logs -f $POD_NAME -c tensorflow

Remarque :L'entraînement du modèle prend environ cinq minutes.