The following section provides an example of fine-tuning a BERT model for sequence classification using the Hugging Face transformers library with TensorFlow. The dataset is downloaded into a mounted Parallelstore-backed volume, allowing the model training to directly read data from the volume.
Prerequisites
- Ensure your node has at least 8 GiB of memory available.
- Create a PersistentVolumeClaim requesting for a Parallelstore-backed volume.
Save the following YAML manifest (parallelstore-csi-job-example.yaml) for your model training Job.
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 Apply the YAML manifest to the cluster.
kubectl apply -f parallelstore-csi-job-example.yaml
Check your data loading and model training progress with the following command:
POD_NAME=$(kubectl get pod | grep 'parallelstore-csi-job-example' | awk '{print $1}') kubectl logs -f $POD_NAME -c tensorflow