بارهای کاری پروفایل PyTorch XLA

بهینه‌سازی عملکرد بخش مهمی از ساخت مدل‌های کارآمد یادگیری ماشین است. شما می‌توانید از ابزار پروفایلینگ XProf برای اندازه‌گیری عملکرد بارهای کاری یادگیری ماشین خود استفاده کنید. XProf به شما امکان می‌دهد تا ردیابی‌های دقیقی از اجرای مدل خود در دستگاه‌های XLA ثبت کنید. این ردیابی‌ها می‌توانند به شما در شناسایی گلوگاه‌های عملکرد، درک میزان استفاده از دستگاه و بهینه‌سازی کد شما کمک کنند.

این راهنما فرآیند ثبت برنامه‌ریزی‌شده‌ی یک رد از اسکریپت PyTorch XLA و مصورسازی آن با استفاده از XProf را شرح می‌دهد.

ضبط یک ردپا به صورت برنامه‌نویسی‌شده

شما می‌توانید با اضافه کردن چند خط کد به اسکریپت آموزشی موجود خود، یک ردیابی را ثبت کنید. ابزار اصلی برای ثبت ردیابی، ماژول torch_xla.debug.profiler است که معمولاً با نام مستعار xp وارد می‌شود.

۱. سرور پروفایلر را شروع کنید

قبل از اینکه بتوانید ردیابی را ضبط کنید، باید سرور پروفایلر را شروع کنید. این سرور در پس‌زمینه اسکریپت شما اجرا می‌شود و داده‌های ردیابی را جمع‌آوری می‌کند. می‌توانید آن را با فراخوانی xp.start_server(<port>) در نزدیکی ابتدای بلوک اجرای اصلی خود شروع کنید.

۲. مدت زمان ردیابی را تعریف کنید

کدی را که می‌خواهید پروفایل کنید، درون فراخوانی‌های xp.start_trace() و xp.stop_trace() قرار دهید. تابع start_trace مسیری را به دایرکتوری که فایل‌های ردیابی در آن ذخیره شده‌اند، دریافت می‌کند.

این یک روش معمول است که حلقه اصلی آموزش را برای ثبت مرتبط‌ترین عملیات، در بر می‌گیرد.

import torch_xla.debug.profiler as xp

# 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()

۳. برچسب‌های ردیابی سفارشی اضافه کنید

به طور پیش‌فرض، ردپاهای ثبت‌شده، توابع سطح پایین 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 به صورت بصری نمایش دهید.

شما می‌توانید رابط کاربری پروفایلر را مستقیماً با استفاده از دستور مستقل XProf و با ارجاع آن به دایرکتوری لاگ خود، اجرا کنید:

$ xprof --port=8791 /root/logs/
Attempting to start XProf server:
  Log Directory: /root/logs/
  Port: 8791
  Worker Service Address: 0.0.0.0:50051
  Hide Capture Button: False
XProf at http://localhost:8791/ (Press CTRL+C to quit)

برای مشاهده پروفایل، به آدرس اینترنتی ارائه شده (مثلاً http://localhost:8791/) در مرورگر خود بروید.

شما قادر خواهید بود برچسب‌های سفارشی که ایجاد کرده‌اید را مشاهده کنید و زمان اجرای بخش‌های مختلف مدل خود را تجزیه و تحلیل کنید.

اگر از Google Cloud برای اجرای بارهای کاری خود استفاده می‌کنید، ابزار cloud-diagnostics-xprof را توصیه می‌کنیم. این ابزار با استفاده از ماشین‌های مجازی که XProf را اجرا می‌کنند، جمع‌آوری و مشاهده پروفایل‌ها را به صورت ساده و روان ارائه می‌دهد.