Vertex AI: אימון רב-עובדים והעברת למידה באמצעות TensorFlow

1. סקירה כללית

בשיעור ה-Lab הזה תשתמשו ב-Vertex AI כדי להריץ משימת אימון של מודל TensorFlow עם כמה עובדים.

מה לומדים

במאמר הזה נסביר איך:

  • שינוי קוד האפליקציה לאימון לצורך אימון עם כמה תהליכי עבודה
  • הגדרה והפעלה של משימת אימון עם כמה עובדים מממשק המשתמש של Vertex AI
  • הגדרה והפעלה של משימת אימון עם כמה עובדים באמצעות Vertex SDK

העלות הכוללת להרצת שיעור ה-Lab הזה ב-Google Cloud היא בערך 5$.

2. מבוא ל-Vertex AI

בשיעור ה-Lab הזה נעשה שימוש במוצר ה-AI החדש ביותר שזמין ב-Google Cloud. ‫Vertex AI משלב את מוצרי ה-ML ב-Google Cloud לחוויית פיתוח חלקה. בעבר, היה אפשר לגשת למודלים שאומנו באמצעות AutoML ולמודלים בהתאמה אישית דרך שירותים נפרדים. המוצר החדש משלב את שניהם ב-API אחד, יחד עם מוצרים חדשים אחרים. אפשר גם להעביר פרויקטים קיימים אל Vertex AI. אם יש לך משוב, אפשר לעיין בדף התמיכה.

‫Vertex AI כולל מוצרים רבים ושונים לתמיכה בתהליכי עבודה של למידת מכונה מקצה לקצה. בשיעור ה-Lab הזה נתמקד במוצרים שמודגשים בהמשך: Training ו-Workbench

סקירה כללית על מוצר Vertex

3. סקירה כללית של תרחיש לדוגמה

בשיעור ה-Lab הזה תשתמשו בלמידת העברה כדי לאמן מודל לסיווג תמונות במערך נתוני הקסאווה מתוך TensorFlow Datasets. הארכיטקטורה שבה תשתמשו היא מודל ResNet50 מהספרייה tf.keras.applications שעבר אימון מראש על מערך הנתונים Imagenet.

למה כדאי להשתמש באימון מבוזר?

אם יש לכם GPU אחד, TensorFlow ישתמש במאיץ הזה כדי להאיץ את אימון המודל בלי שתצטרכו לעשות שום דבר נוסף. עם זאת, אם רוצים לקבל דחיפה נוספת מהשימוש בכמה יחידות GPU במכונה אחת או בכמה מכונות (כל אחת עם כמה יחידות GPU פוטנציאליות), צריך להשתמש ב-tf.distribute, שהיא ספריית TensorFlow להפעלת חישוב בכמה מכשירים. מכשיר הוא מעבד (CPU) או מאיץ, כמו GPUs או TPUs, במכונה מסוימת ש-TensorFlow יכולה להריץ עליה פעולות.

הדרך הכי פשוטה להתחיל באימון מבוזר היא שימוש במכונה אחת עם כמה מכשירי GPU. אסטרטגיית הפצה של TensorFlow מהמודול tf.distribute תנהל את התיאום של הפצת הנתונים ועדכוני הגרדיאנט בכל יחידות ה-GPU. אם אתם שולטים באימון של מארח יחיד ומחפשים דרכים להרחבה נוספת, הוספה של כמה מכונות לאשכול יכולה לעזור לכם לשפר עוד יותר את הביצועים. אפשר להשתמש באוסף של מכונות שיש בהן רק מעבדים מרכזיים, או שבכל אחת מהן יש מעבד גרפי אחד או יותר. בשיעור ה-Lab הזה נתמקד במקרה השני ונדגים איך להשתמש ב-MultiWorkerMirroredStrategy כדי להפיץ את האימון של מודל TensorFlow בין כמה מכונות ב-Vertex AI.

‫MultiWorkerMirroredStrategy היא אסטרטגיה סינכרונית של מקביליות נתונים שאפשר להשתמש בה עם כמה שינויים קלים בקוד. עותק של המודל נוצר בכל מכשיר באשכול. העדכונים הבאים של הגרדיאנט יתבצעו באופן סינכרוני. כלומר, כל מכשיר עובד מחשב את המעברים קדימה ואחורה דרך המודל על פרוסת נתונים שונה של נתוני הקלט. הגרדיאנטים המחושבים מכל אחד מהפלחים האלה מצטברים בכל המכשירים במכונה ובכל המכונות באשכול, ועוברים צמצום (בדרך כלל ממוצע) בתהליך שנקרא all-reduce. לאחר מכן, האופטימיזציה מבצעת את עדכוני הפרמטרים עם הגרדיאנטים המופחתים האלה, וכך המכשירים נשארים מסונכרנים. מידע נוסף על אימון מבוזר באמצעות TensorFlow זמין בסרטון הבא:

4. הגדרת הסביבה

כדי להפעיל את ה-codelab הזה, צריך פרויקט ב-Google Cloud Platform שמופעל בו חיוב. כדי ליצור פרויקט, פועלים לפי ההוראות האלה.

שלב 1: הפעלת Compute Engine API

עוברים אל Compute Engine ובוחרים באפשרות הפעלה אם הוא עדיין לא מופעל. תצטרכו את זה כדי ליצור את מופע המחברת.

שלב 2: הפעלת Container Registry API

עוברים אל Container Registry ובוחרים באפשרות Enable (הפעלה) אם היא עדיין לא מסומנת. תשתמשו בזה כדי ליצור קונטיינר למשימת האימון המותאמת אישית.

שלב 3: הפעלת Vertex AI API

עוברים אל הקטע Vertex AI במסוף Cloud ולוחצים על הפעלת Vertex AI API.

לוח הבקרה של Vertex AI

שלב 4: יצירת מכונה של Vertex AI Workbench

בקטע Vertex AI במסוף Cloud, לוחצים על Workbench:

תפריט Vertex AI

מפעילים את Notebooks API אם הוא עדיין לא מופעל.

Notebook_api

אחרי ההפעלה, לוחצים על מחברות מנוהלות:

Notebooks_UI

לאחר מכן בוחרים באפשרות מחברת חדשה.

new_notebook

נותנים שם למחברת ולוחצים על הגדרות מתקדמות.

create_notebook

בקטע 'הגדרות מתקדמות', מפעילים את ההגדרה 'כיבוי במצב לא פעיל' ומגדירים את מספר הדקות ל-60. המשמעות היא שמחברת ה-notebook תיסגר אוטומטית כשלא משתמשים בה, כדי שלא תצטברו עלויות מיותרות.

idle_timeout

בקטע אבטחה, בוחרים באפשרות 'הפעלת מסוף' אם היא עדיין לא מופעלת.

enable_terminal

אפשר להשאיר את כל ההגדרות המתקדמות האחרות כמו שהן.

אחרי כן, לוחצים על יצירה. ייקח כמה דקות עד שהמופע יוקצה.

אחרי שהמופע נוצר, בוחרים באפשרות Open JupyterLab.

open_jupyterlab

בפעם הראשונה שמשתמשים במופע חדש, תופיע בקשה לבצע אימות. פועלים לפי השלבים בממשק המשתמש.

אימות

5. יצירת קונטיינר לקוד אפליקציה של אימון

כדי לשלוח את משימת האימון הזו ל-Vertex, צריך להכניס את קוד אפליקציית האימון למאגר Docker ולדחוף את המאגר הזה ל-Google Container Registry. באמצעות הגישה הזו, אפשר לאמן מודל שנבנה עם כל מסגרת.

כדי להתחיל, בתפריט מרכז האפליקציות, פותחים חלון הטרמינל במופע של מחברת:

פתיחת טרמינל ב-notebook

יוצרים ספרייה חדשה בשם cassava ועוברים אליה באמצעות הפקודה cd:

mkdir cassava
cd cassava

שלב 1: יצירת Dockerfile

השלב הראשון בהעברת הקוד למאגר הוא יצירת Dockerfile. ב-Dockerfile, צריך לכלול את כל הפקודות שנדרשות להרצת האימג'. הסקריפט יתקין את כל הספריות הנדרשות ויגדיר את נקודת הכניסה לקוד האימון.

ב-Terminal, יוצרים קובץ Dockerfile ריק:

touch Dockerfile

פותחים את Dockerfile ומעתיקים לתוכו את הקוד הבא:

FROM gcr.io/deeplearning-platform-release/tf2-gpu.2-7

WORKDIR /

# Copies the trainer code to the docker image.
COPY trainer /trainer

# Sets up the entry point to invoke the trainer.
ENTRYPOINT ["python", "-m", "trainer.task"]

קובץ ה-Dockerfile הזה משתמש בקובץ אימג' של Docker ל-GPU של TensorFlow Enterprise 2.7 של Deep Learning Container. הקונטיינרים של Deep Learning ב-Google Cloud מגיעים עם הרבה frameworks נפוצים של ML ומדעי נתונים שהותקנו מראש. אחרי שמורידים את התמונה הזו, קובץ ה-Docker הזה מגדיר את נקודת הכניסה לקוד האימון. עדיין לא יצרתם את הקבצים האלה – בשלב הבא תוסיפו את הקוד לאימון המודל ולכוונון שלו.

שלב 2: יצירת קטגוריה של Cloud Storage

במשימת האימון הזו תייצאו את מודל TensorFlow שאומן לקטגוריה של Cloud Storage. במסוף, מריצים את הפקודה הבאה כדי להגדיר משתנה סביבה לפרויקט. חשוב להחליף את your-cloud-project במזהה הפרויקט:

PROJECT_ID='your-cloud-project'

לאחר מכן, מריצים את הפקודה הבאה בטרמינל כדי ליצור קטגוריה חדשה בפרויקט.

BUCKET="gs://${PROJECT_ID}-bucket"
gsutil mb -l us-central1 $BUCKET

שלב 3: הוספת קוד לאימון המודל

במסוף, מריצים את הפקודה הבאה כדי ליצור ספרייה לקוד האימון וקובץ Python שבו תוסיפו את הקוד:

mkdir trainer
touch trainer/task.py

עכשיו אמורים להיות לכם הקבצים הבאים בספרייה cassava/:

+ Dockerfile
+ trainer/
    + task.py

לאחר מכן, פותחים את הקובץ task.py שיצרתם ומעתיקים את הקוד שבהמשך. צריך להחליף את {your-gcs-bucket} בשם של הקטגוריה של Cloud Storage שיצרתם.

import tensorflow as tf
import tensorflow_datasets as tfds
import os


PER_REPLICA_BATCH_SIZE = 64
EPOCHS = 2

# TODO: replace {your-gcs-bucket} with the name of the Storage bucket you created earlier
BUCKET = 'gs://{your-gcs-bucket}/mwms'

def preprocess_data(image, label):
  '''Resizes and scales images.'''

  image = tf.image.resize(image, (300,300))
  return tf.cast(image, tf.float32) / 255., label


def create_dataset(batch_size):
  '''Loads Cassava dataset and preprocesses data.'''

  data, info = tfds.load(name='cassava', as_supervised=True, with_info=True)
  number_of_classes = info.features['label'].num_classes
  train_data = data['train'].map(preprocess_data,
                                 num_parallel_calls=tf.data.experimental.AUTOTUNE)
  train_data  = train_data.shuffle(1000)
  train_data  = train_data.batch(batch_size)
  train_data  = train_data.prefetch(tf.data.experimental.AUTOTUNE)

  # Set AutoShardPolicy
  options = tf.data.Options()
  options.experimental_distribute.auto_shard_policy = tf.data.experimental.AutoShardPolicy.DATA
  train_data = train_data.with_options(options)

  return train_data, number_of_classes


def create_model(number_of_classes):
  '''Creates and compiles pretrained ResNet50 model.'''

  base_model = tf.keras.applications.ResNet50(weights='imagenet', include_top=False)
  x = base_model.output
  x = tf.keras.layers.GlobalAveragePooling2D()(x)
  x = tf.keras.layers.Dense(1016, activation='relu')(x)
  predictions = tf.keras.layers.Dense(number_of_classes, activation='softmax')(x)
  model = tf.keras.Model(inputs=base_model.input, outputs=predictions)

  model.compile(
      loss='sparse_categorical_crossentropy',
      optimizer=tf.keras.optimizers.Adam(0.0001),
      metrics=['accuracy'])

  return model


def _is_chief(task_type, task_id):
  '''Helper function. Determines if machine is chief.'''

  return task_type == 'chief'


def _get_temp_dir(dirpath, task_id):
  '''Helper function. Gets temporary directory for saving model.'''

  base_dirpath = 'workertemp_' + str(task_id)
  temp_dir = os.path.join(dirpath, base_dirpath)
  tf.io.gfile.makedirs(temp_dir)
  return temp_dir


def write_filepath(filepath, task_type, task_id):
  '''Helper function. Gets filepath to save model.'''

  dirpath = os.path.dirname(filepath)
  base = os.path.basename(filepath)
  if not _is_chief(task_type, task_id):
    dirpath = _get_temp_dir(dirpath, task_id)
  return os.path.join(dirpath, base)


def main():
  # Create strategy
  strategy = tf.distribute.MultiWorkerMirroredStrategy()

  # Get data
  global_batch_size = PER_REPLICA_BATCH_SIZE * strategy.num_replicas_in_sync
  train_data, number_of_classes = create_dataset(global_batch_size)

  # Wrap variable creation within strategy scope
  with strategy.scope():
    model = create_model(number_of_classes)

  model.fit(train_data, epochs=EPOCHS)

  # Determine type and task of the machine from
  # the strategy cluster resolver
  task_type, task_id = (strategy.cluster_resolver.task_type,
                        strategy.cluster_resolver.task_id)

  # Based on the type and task, write to the desired model path
  write_model_path = write_filepath(BUCKET, task_type, task_id)
  model.save(write_model_path)

if __name__ == "__main__":
    main()

לפני שיוצרים את מאגר התגים, כדאי לעיין בקוד שמשתמש ב-MultiWorkerMirroredStrategy מ-API‏ tf.distribute.Strategy.

יש כמה רכיבים בקוד שנדרשים כדי שהקוד יפעל עם MultiWorkerMirroredStrategy.

  1. הנתונים צריכים להיות מחולקים, כלומר לכל עובד מוקצית קבוצת משנה של מערך הנתונים כולו. לכן, בכל שלב, כל עובד יעבד גודל אצווה גלובלי של רכיבי מערך נתונים שלא חופפים. החלוקה הזו מתבצעת אוטומטית באמצעות tf.data.experimental.AutoShardPolicy, שאפשר להגדיר לו את הערכים FILE או DATA. בדוגמה הזו, הפונקציה create_dataset() מגדירה את AutoShardPolicy ל-DATA כי מערך הנתונים של הקסאווה לא מורד כמספר קבצים. עם זאת, אם לא הגדרתם את המדיניות ל-DATA, מדיניות ברירת המחדל AUTO תיכנס לתוקף והתוצאה הסופית תהיה זהה. כאן אפשר לקרוא מידע נוסף על חלוקת מערכי נתונים באמצעות MultiWorkerMirroredStrategy.
  2. האובייקט MultiWorkerMirroredStrategy נוצר בפונקציה main(). בשלב הבא, עוטפים את יצירת משתני המודל בהיקף של שיטת הבידינג. השלב הזה חשוב מאוד כי הוא קובע ב-TensorFlow אילו משתנים יש לשכפל בכל הרפליקות.
  3. גודל האצווה גדל פי num_replicas_in_sync. כך מוודאים שכל רפליקה מעבדת את אותו מספר של דוגמאות בכל שלב. הגדלת גודל האצווה היא שיטה מומלצת כשמשתמשים בשיטות מקביליות סינכרוניות של נתונים ב-TensorFlow.
  4. שמירת המודל במקרה של כמה עובדים היא קצת יותר מורכבת, כי היעד צריך להיות שונה לכל אחד מהעובדים. התהליך המרכזי ישמור את המודל בספריית המודלים הרצויה, ותהליכים אחרים ישמרו את המודל בספריות זמניות. חשוב שהספריות הזמניות האלה יהיו ייחודיות כדי למנוע ממספר עובדים לכתוב לאותו מיקום. השמירה יכולה לכלול פעולות משותפות, כלומר כל העובדים צריכים לשמור ולא רק המנהל. הפונקציות _is_chief(), ‏ _get_temp_dir(), ‏write_filepath(), וגם הפונקציה main() כוללות קוד boilerplate שעוזר לשמור את המודל.

שימו לב: אם השתמשתם ב-MultiWorkerMirroredStrategy בסביבה אחרת, יכול להיות שהגדרתם את משתנה הסביבה TF_CONFIG. ‫Vertex AI מגדיר את TF_CONFIG באופן אוטומטי, כך שלא צריך להגדיר את המשתנה הזה בכל מכונה באשכול.

שלב 4: בניית הקונטיינר

במסוף, מריצים את הפקודה הבאה כדי להגדיר משתנה סביבה לפרויקט. חשוב להחליף את your-cloud-project במזהה הפרויקט:

PROJECT_ID='your-cloud-project'

מגדירים משתנה עם ה-URI של קובץ אימג' של קונטיינר ב-Google Container Registry:

IMAGE_URI="gcr.io/$PROJECT_ID/multiworker:cassava"

הגדרת Docker

gcloud auth configure-docker

לאחר מכן, בונים את הקונטיינר על ידי הרצת הפקודה הבאה מתיקיית השורש של cassava:

docker build ./ -t $IMAGE_URI

לבסוף, מעלים אותו אל Google Container Registry:

docker push $IMAGE_URI

אחרי שמעבירים את המאגר אל Container Registry, אפשר להפעיל את משימת האימון.

6. הרצה של משימת אימון עם כמה workers ב-Vertex AI

במעבדה הזו נעשה שימוש באימון בהתאמה אישית באמצעות קונטיינר בהתאמה אישית ב-Google Container Registry, אבל אפשר גם להריץ משימת אימון באמצעות קונטיינרים מוכנים מראש.

כדי להתחיל, עוברים לקטע Training בקטע Vertex במסוף Cloud:

תפריט uCAIP

שלב 1: הגדרת משימת אימון

לוחצים על יצירה כדי להזין את הפרמטרים של משימת האימון.

  • בקטע מערך נתונים, בוחרים באפשרות אין מערך נתונים מנוהל.
  • לאחר מכן בוחרים באפשרות Custom training (advanced) (אימון בהתאמה אישית (מתקדם)) בתור שיטת האימון ולוחצים על Continue (המשך).
  • מזינים multiworker-cassava (או כל שם אחר שרוצים לתת למודל) בשדה שם המודל.
  • לוחצים על המשך.

בשלב Container settings (הגדרות מאגר התגים), בוחרים באפשרות Custom container (מאגר תגים בהתאמה אישית):

אפשרות של מאגר תגים מותאם אישית

בתיבה הראשונה (קובץ אימג' של קונטיינר), מזינים את הערך של המשתנה IMAGE_URI מהקטע הקודם. הפורמט צריך להיות: gcr.io/your-cloud-project/multiworker:cassava, עם מזהה הפרויקט שלכם. משאירים את שאר השדות ריקים ולוחצים על המשך.

מדלגים על שלב ההיפרפרמטרים בלחיצה נוספת על המשך.

שלב 2: הגדרת אשכול מחשוב

‫Vertex AI מספק 4 מאגרי עובדים כדי לכסות את הסוגים השונים של משימות למידת מכונה.

Worker pool 0 מגדיר את ה-Primary,‏ chief,‏ scheduler או master. ב-MultiWorkerMirroredStrategy, כל המכונות מוגדרות כעובדות, שהן המכונות הפיזיות שבהן מבוצע החישוב המשוכפל. בנוסף לכך שכל מכונה היא worker, צריך להיות worker אחד שמבצע עבודה נוספת, כמו שמירת נקודות ביקורת וכתיבת קובצי סיכום ל-TensorBoard. המכונה הזו נקראת 'הבוס'. תמיד יש רק עובד ראשי אחד, לכן מספר העובדים במאגר העובדים 0 תמיד יהיה 1.

בקטע Compute and pricing (חישוב ותמחור), משאירים את האזור שנבחר כמו שהוא ומגדירים את Worker pool 0 (מאגר עובדים 0) באופן הבא:

Worker_pool_0

במאגר העובדים 1 מגדירים את העובדים עבור האשכול.

מגדירים את מאגר העובדים 1 באופן הבא:

Worker_pool_1

האשכול מוגדר עכשיו עם שני מכשירים עם מעבד בלבד. כשמריצים את קוד אפליקציית האימון, MultiWorkerMirroredStrategy יפיץ את האימון על פני שני המחשבים.

MultiWorkerMirroredStrategy יש רק את סוגי המשימות chief ו-worker, כך שאין צורך להגדיר את מאגרי העובדים הנוספים. עם זאת, אם הייתם משתמשים ב-ParameterServerStrategy של TensorFlow, הייתם מגדירים את שרתי הפרמטרים במאגר העובדים 2. אם רוצים להוסיף מכונת הערכה לאשכול, צריך להגדיר את המכונה במאגר העובדים 3.

כדי להפעיל את עבודת הכוונון של ההיפרפרמטרים, לוחצים על Start training (התחלת האימון). בקטע Training במסוף, בכרטיסייה TRAINING PIPELINES, תופיע המשימה שהופעלה:

משימות אימון

‫🎉 איזה כיף! 🎉

למדתם איך להשתמש ב-Vertex AI כדי:

  • הפעלת משימת אימון עם כמה worker לאימון קוד שסופק בקונטיינר מותאם אישית. בדוגמה הזו השתמשתם במודל TensorFlow, אבל אתם יכולים לאמן מודל שנבנה עם כל מסגרת באמצעות קונטיינרים מובנים או קונטיינרים בהתאמה אישית.

מידע נוסף על החלקים השונים של Vertex זמין בתיעוד.

7. [אופציונלי] שימוש ב-Vertex SDK

בקטע הקודם הראינו איך להפעיל את עבודת האימון דרך ממשק המשתמש. בקטע הזה מוצגת דרך חלופית לשליחת משימת האימון באמצעות API בשפת Python של Vertex.

חזרו למופע של המחברת וצרו מחברת TensorFlow 2 מ-Launcher:

new_notebook

מייבאים את Vertex AI SDK.

from google.cloud import aiplatform

כדי להפעיל את משימת האימון עם כמה עובדים, צריך קודם להגדיר את המפרט של מאגר העובדים. שימו לב שהשימוש ב-GPU במפרט הוא אופציונלי לחלוטין, ואפשר להסיר את accelerator_type ואת accelerator_count אם רוצים להשתמש באשכול עם CPU בלבד, כמו שמוצג בקטע הקודם.

# The spec of the worker pools including machine type and Docker image
# Be sure to replace {YOUR-PROJECT-ID} with your project ID.
worker_pool_specs=[
     {
        "replica_count": 1,
        "machine_spec": {
          "machine_type": "n1-standard-8", "accelerator_type": "NVIDIA_TESLA_V100", "accelerator_count": 1
        },
        "container_spec": {"image_uri": "gcr.io/{YOUR-PROJECT-ID}/multiworker:cassava"}
      },
      {
        "replica_count": 1,
        "machine_spec": {
          "machine_type": "n1-standard-8", "accelerator_type": "NVIDIA_TESLA_V100", "accelerator_count": 1
        },
        "container_spec": {"image_uri": "gcr.io/{YOUR-PROJECT-ID}/multiworker:cassava"}
      }
]

לאחר מכן, יוצרים ומריצים CustomJob. צריך להחליף את {YOUR_BUCKET} בקטגוריה בפרויקט שלכם לצורך העברה זמנית. אתם יכולים להשתמש באותה קטגוריה שיצרתם קודם.

# Replace YOUR_BUCKET
my_multiworker_job = aiplatform.CustomJob(display_name='multiworker-cassava-sdk',
                              worker_pool_specs=worker_pool_specs,
                              staging_bucket='gs://{YOUR_BUCKET}')

my_multiworker_job.run()

בקטע Training (אימון) במסוף, בכרטיסייה CUSTOM JOBS (משימות בהתאמה אישית), תופיע משימת האימון:

משרות בהתאמה אישית

8. הסרת המשאבים

הגדרנו את מחברת ה-Jupyter כך שתפסיק לפעול אחרי 60 דקות של חוסר פעילות, ולכן אין צורך לדאוג להשבתת המופע. כדי לכבות את המופע באופן ידני, לוחצים על הלחצן Stop (עצירה) בקטע Vertex AI Workbench במסוף. כדי למחוק את ה-Notebook לגמרי, לוחצים על לחצן המחיקה.

עצירת מכונה

כדי למחוק את קטגוריית האחסון, בתפריט הניווט ב-Cloud Console, עוברים אל Storage, בוחרים את הקטגוריה ולוחצים על Delete:

מחיקת האחסון