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

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

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

תיעוד מעקב

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

1. הפעלת שרת ה-profiler

כדי לתעד מעקב, צריך להפעיל את שרת ה-profiler. השרת הזה פועל ברקע של הסקריפט ואוסף את נתוני המעקב. כדי להפעיל אותו, אפשר להפעיל את 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. הוא מספק חוויית איסוף וצפייה יעילה בפרופילים באמצעות מכונות וירטואליות שמריצות את Tensorboard ו-XProf.