परफ़ॉर्मेंस ऑप्टिमाइज़ेशन, असरदार मशीन लर्निंग मॉडल बनाने का एक अहम हिस्सा है. मशीन लर्निंग के वर्कलोड की परफ़ॉर्मेंस का आकलन करने के लिए, 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() कॉन्टेक्स्ट मैनेजर का इस्तेमाल करके, अपने कोड के कुछ सेक्शन में कस्टम लेबल जोड़े जा सकते हैं. ये लेबल, प्रोफ़ाइलर के टाइमलाइन व्यू में नाम वाले ब्लॉक के तौर पर दिखेंगे. इससे डेटा तैयार करने, फ़ॉरवर्ड पास या ऑप्टिमाइज़र चरण जैसे खास ऑपरेशन की पहचान करना बहुत आसान हो जाएगा.
यहां दिए गए उदाहरण में, ट्रेनिंग के चरण के अलग-अलग हिस्सों में कॉन्टेक्स्ट जोड़ने का तरीका बताया गया है.
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()
पूरा उदाहरण
यहां दिए गए उदाहरण में, mnist_xla.py फ़ाइल के आधार पर, PyTorch XLA स्क्रिप्ट से ट्रेस कैप्चर करने का तरीका बताया गया है.
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 चलाने वाले वीएम का इस्तेमाल करके, प्रोफ़ाइल इकट्ठा करने और देखने का आसान अनुभव देता है.