פרופיל של עומסי עבודה ב-PyTorch XLA

אופטימיזציה של הביצועים היא חלק חשוב בבניית מודלים יעילים של למידת מכונה. אתם יכולים להשתמש בכלי ליצירת פרופילים XProf כדי למדוד את הביצועים של עומסי העבודה של למידת המכונה. ‫XProf מאפשר לכם ללכוד עקבות מפורטים של ההרצה של המודל במכשירי XLA. הנתונים האלה יכולים לעזור לכם לזהות צווארי בקבוק בביצועים, להבין את השימוש במכשיר ולבצע אופטימיזציה של הקוד.

במדריך הזה מוסבר תהליך הלכידה הפרוגרמטית של נתוני מעקב מסקריפט PyTorch XLA והצגתם באמצעות XProf ו-Tensorboard.

תיעוד עקבות

כדי לתעד את הנתונים, מוסיפים כמה שורות קוד לסקריפט האימון הקיים. הכלי העיקרי ללכידת נתוני מעקב הוא המודול torch_xla.debug.profiler, שבדרך כלל מיובא עם הכינוי xp.

1. הפעלת שרת הפרופיל

כדי ללכוד נתונים למעקב, צריך להפעיל את שרת הפרופיל. השרת הזה פועל ברקע של הסקריפט ואוסף את נתוני המעקב. אפשר להפעיל אותו על ידי קריאה ל-xp.start_server() קרוב לתחילת בלוק הביצוע הראשי.

2. הגדרת משך המעקב

עוטפים את הקוד שרוצים ליצור לו פרופיל בקריאות xp.start_trace() ו-xp.stop_trace(). הפונקציה start_trace מקבלת נתיב לספרייה שבה נשמרים קובצי המעקב.

נהוג לעטוף את לולאת האימון הראשית כדי לתעד את הפעולות הרלוונטיות ביותר.

# The directory where the trace files are stored.
log_dir = '/root/logs/'

# Start tracing
xp.start_trace(log_dir)

# ... your training loop or other code to be profiled ...
train_mnist()

# Stop tracing
xp.stop_trace()

3. הוספת תוויות מותאמות אישית למעקב

כברירת מחדל, העקבות שמתועדים הן פונקציות ברמה נמוכה של Pytorch XLA, ולפעמים קשה לנווט בהן. אפשר להוסיף תוויות מותאמות אישית לקטעים ספציפיים בקוד באמצעות xp.Trace() מנהל ההקשר. התוויות האלה יופיעו כבלוקים עם שמות בתצוגת ציר הזמן של כלי הפרופיל, וכך יהיה קל יותר לזהות פעולות ספציפיות כמו הכנת נתונים, העברה קדימה או שלב האופטימיזציה.

בדוגמה הבאה אפשר לראות איך מוסיפים הקשר לחלקים שונים בשלב ההכשרה.

def forward(self, x):
    # This entire block will be labeled 'forward' in the trace
    with xp.Trace('forward'):
        x = F.relu(F.max_pool2d(self.conv1(x), 2))
        x = F.relu(F.max_pool2d(self.conv2(x), 2))
        x = x.view(-1, 7*7*64)
        x = F.relu(self.fc1(x))
        x = self.fc2(x)
        return F.log_softmax(x, dim=1)

# You can also nest context managers for more granular detail
for batch_idx, (data, target) in enumerate(train_loader):
    with torch_xla.step():
        with xp.Trace('train_step_data_prep_and_forward'):
            optimizer.zero_grad()
            data, target = data.to(device), target.to(device)
            output = model(data)

        with xp.Trace('train_step_loss_and_backward'):
            loss = loss_fn(output, target)
            loss.backward()

        with xp.Trace('train_step_optimizer_step_host'):
            optimizer.step()

דוגמה מלאה

בדוגמה הבאה אפשר לראות איך ללכוד נתונים למעקב מסקריפט PyTorch XLA, על סמך הקובץ mnist_xla.py.

import torch
import torch.optim as optim
from torchvision import datasets, transforms

# PyTorch/XLA specific imports
import torch_xla
import torch_xla.core.xla_model as xm
import torch_xla.debug.profiler as xp

def train_mnist():
    # ... (model definition and data loading code) ...
    print("Starting training...")
    # ... (training loop as defined in the previous section) ...
    print("Training finished!")

if __name__ == '__main__':
    # 1. Start the profiler server
    server = xp.start_server(9012)

    # 2. Start capturing the trace and define the output directory
    xp.start_trace('/root/logs/')

    # Run the training function that contains custom trace labels
    train_mnist()

    # 3. Stop the trace
    xp.stop_trace()

הדמיה של העקבות

כשהסקריפט מסתיים, קובצי המעקב נשמרים בספרייה שציינתם (לדוגמה, /root/logs/). אפשר להציג את המעקב הזה באמצעות XProf ו-TensorBoard.

  1. מתקינים את TensorBoard.

    pip install tensorboard_plugin_profile tensorboard

  2. מפעילים את TensorBoard. מפנים את TensorBoard לספריית היומן שבה השתמשתם ב-xp.start_trace():

    tensorboard --logdir /root/logs/

  3. הצגת הפרופיל. פותחים את כתובת ה-URL שסופקה על ידי TensorBoard בדפדפן האינטרנט (בדרך כלל http://localhost:6006). עוברים לכרטיסייה PyTorch XLA - Profile כדי לראות את המעקב האינטראקטיבי. תוכלו לראות את התוויות המותאמות אישית שיצרתם ולנתח את זמן הביצוע של חלקים שונים במודל.

אם אתם משתמשים ב- Google Cloud להפעלת עומסי העבודה, מומלץ להשתמש בכלי cloud-diagnostics-xprof. הוא מספק חוויה יעילה של איסוף פרופילים וצפייה בהם באמצעות מכונות וירטואליות (VM) שמריצות Tensorboard ו-XProf.