ملفات تعريف أحمال عمل PyTorch XLA

يُعدّ تحسين الأداء جزءًا مهمًا من إنشاء نماذج فعّالة لتعلُّم الآلة. يمكنك استخدام أداة XProf لإنشاء الملفات الشخصية من أجل قياس أداء أحمال عمل تعلُّم الآلة. تتيح لك أداة XProf تسجيل عمليات تتبُّع مفصّلة لتنفيذ النموذج على أجهزة XLA. يمكن أن تساعدك عمليات التتبُّع هذه في تحديد المؤثِّرات السلبية في الأداء وفهم معدّل استخدام الجهاز وتحسين الرمز البرمجي.

يوضّح هذا الدليل عملية تسجيل عملية تتبُّع برمجيًا من نص PyTorch XLA البرمجي وعرضها باستخدام XProf.

تسجيل عملية تتبُّع آليًا

يمكنك تسجيل عملية تتبُّع من خلال إضافة بضعة أسطر من الرمز إلى البرنامج النصي الحالي الخاص بالتدريب. الأداة الأساسية لتسجيل التتبُّع هي الوحدة torch_xla.debug.profiler، والتي يتم استيرادها عادةً بالاسم المستعار xp.

1- بدء خادم أداة تحليل الأداء

قبل أن تتمكّن من تسجيل تتبُّع، عليك بدء خادم أداة تحليل الأداء. يعمل هذا الخادم في خلفية النص البرمجي ويجمع بيانات التتبُّع. يمكنك بدء عملية التتبُّع من خلال استدعاء xp.start_server(<port>) بالقرب من بداية كتلة التنفيذ الرئيسية.

2. تحديد مدة التتبُّع

ضَع الرمز الذي تريد إنشاء ملف تعريف له بين استدعاءات 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()

3- إضافة تصنيفات مخصّصة لعمليات التتبُّع

تكون عمليات التتبُّع التي يتم تسجيلها تلقائيًا عبارة عن دوال Pytorch XLA منخفضة المستوى، وقد يكون من الصعب التنقّل فيها. يمكنك إضافة تصنيفات مخصّصة إلى أقسام معيّنة من الرمز باستخدام أداة إدارة السياق xp.Trace(). ستظهر هذه التصنيفات على شكل مربّعات مسماة في عرض المخطط الزمني لأداة Profiler، ما يسهّل كثيرًا تحديد عمليات معيّنة، مثل إعداد البيانات أو التمرير الأمامي أو خطوة المحسِّن.

يوضّح المثال التالي كيف يمكنك إضافة سياق إلى أجزاء مختلفة من خطوة تدريبية.

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.

يمكنك تشغيل واجهة مستخدم أداة Profiler مباشرةً باستخدام أمر 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)

انتقِل إلى عنوان URL المقدَّم (مثل http://localhost:8791/) في المتصفّح لعرض الملف الشخصي.

ستتمكّن من الاطّلاع على التصنيفات المخصّصة التي أنشأتها وتحليل وقت التنفيذ لأجزاء مختلفة من نموذجك.

إذا كنت تستخدم Google Cloud لتشغيل أحمال العمل، ننصحك باستخدام أداة cloud-diagnostics-xprof. ويوفّر تجربة مبسطة لجمع الملفات الشخصية وعرضها باستخدام الآلات الافتراضية التي تشغّل XProf.