PyTorch XLA 워크로드 프로파일링

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

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

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