Pytorch的编译新特性TorchDynamo的工作原理和使用示例

在深度学习中，优化模型性能至关重要，特别是对于需要快速执行和实时推断的应用。而PyTorch在平衡动态图执行与高性能方面常常面临挑战。传统的PyTorch优化技术在处理动态计算图时效果有限，导致训练时间延长和模型性能不佳。TorchDynamo是一种为PyTorch设计的即时（JIT）编译器，通过在运行时拦截Python代码、优化它，并编译成高效的机器代码来解决这一问题。本文通过使用合成数据集展示了TorchDynamo的实际应用，包括特征工程、超参数调整、交叉验证和评估指标。

TorchDynamo简介

TorchDynamo 是一个由 PyTorch 团队开发的编译器前端，它旨在自动优化 PyTorch 程序以提高运行效率。TorchDynamo 的工作原理是在运行时动态分析和转换 PyTorch 的代码，然后将其转发给各种后端编译器（如 TorchScript、TVM、Triton 等），从而实现性能的提升。

特别是在需要实时执行的应用中，如自动驾驶或金融预测等，深度学习模型要求快速执行。传统的优化技术经常需要在处理Python的动态特性时进行修订，这正是TorchDynamo的强项所在。它能够即时捕获计算图，针对特定的工作负载和硬件应用优化，从而减少延迟并提高吞吐量。

TorchDynamo的另外一个突出特点是其易于集成。重写整个代码库以集成新工具可能是一项艰巨的任务。但是TorchDynamo仅需要对现有的PyTorch工作流进行最小的更改。它的简单性和强大的优化能力使它成为经验丰富的研究人员和行业专业人士的有力选择。

将 TorchDynamo 集成到现有的 PyTorch 程序中相对简单，只需要在程序中导入 TorchDynamo 并使用它来包装模型的执行部分。

import torch
import torchdynamo
# 定义模型和优化器
model = MyModel()
optimizer = torch.optim.Adam(model.parameters())
# 使用 TorchDynamo 优化模型的训练过程
def training_step(input, target):
optimizer.zero_grad()
output = model(input)
loss = loss_fn(output, target)
loss.backward()
optimizer.step()
return loss
# 使用 torchdynamo.optimize 包装训练步骤
optimized_training_step = torchdynamo.optimize(training_step)
# 训练循环
for input, target in data_loader:
loss = optimized_training_step(input, target)

TorchDynamo的工作原理

TorchDynamo通过追踪PyTorch代码的执行，动态地捕获计算图。这个过程涉及理解代码的依赖关系和流程，使其能够识别优化的机会。 应用优化

一旦捕获了计算图，TorchDynamo就会应用各种优化技术。这些技术包括操作符融合，它将多个操作合并为一个单一操作以减少开销，以及改进内存管理，最小化数据移动并有效地重用资源。

优化计算图口，TorchDynamo将其编译成高效的机器码。这种编译可以针对不同的后端，如TorchScript或NVFuser，以确保代码在可用的硬件上以最佳方式运行。

在最后的执行阶段。与最初的Python代码相比，上面的优化可以显著提高性能。JIT编译确保在运行时期间应用这些优化，使执行适应不同的工作负载和输入数据。

使用示例

下面我们演示了使用一个合成数据集的TorchDynamo示例，包括特征工程，超参数调优，交叉验证，预测和结果解释。

import torch
import torch.nn as nn
import torch.optim as optim
import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
from sklearn.model_selection import train_test_split, KFold
from sklearn.metrics import mean_squared_error, r2_score
from sklearn.preprocessing import StandardScaler
from torch import _dynamo as torchdynamo
from typing import List
# Generate synthetic dataset
np.random.seed(42)
torch.manual_seed(42)
# Feature engineering: create synthetic data
n_samples = 1000
n_features = 10
X = np.random.rand(n_samples, n_features)
y = X @ np.random.rand(n_features) + np.random.rand(n_samples) * 0.1 # Linear relation with noise
# Split data into train and test sets
X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.2, random_state=42)
# Standardize the features
scaler = StandardScaler()
X_train = scaler.fit_transform(X_train)
X_test = scaler.transform(X_test)
# Convert to PyTorch tensors
X_train = torch.tensor(X_train, dtype=torch.float32)
y_train = torch.tensor(y_train, dtype=torch.float32).view(-1, 1)
X_test = torch.tensor(X_test, dtype=torch.float32)
y_test = torch.tensor(y_test, dtype=torch.float32).view(-1, 1)
# Define the model
class SimpleNN(nn.Module):
def __init__(self, input_dim):
super(SimpleNN, self).__init__()
self.fc1 = nn.Linear(input_dim, 64)
self.fc2 = nn.Linear(64, 32)
self.fc3 = nn.Linear(32, 1)
def forward(self, x):
x = torch.relu(self.fc1(x))
x = torch.relu(self.fc2(x))
x = self.fc3(x)
return x
# Hyperparameters
input_dim = X_train.shape[1]
learning_rate = 0.001
n_epochs = 100
# Initialize the model, loss function, and optimizer
model = SimpleNN(input_dim)
criterion = nn.MSELoss()
optimizer = optim.Adam(model.parameters(), lr=learning_rate)
# Define custom compiler
def my_compiler(gm: torch.fx.GraphModule, example_inputs: List[torch.Tensor]):
print("my_compiler() called with FX graph:")
gm.graph.print_tabular()
return gm.forward # return a python callable
@torchdynamo.optimize(my_compiler)
def train_and_evaluate(model, criterion, optimizer, X_train, y_train, X_test, y_test, n_epochs):
# Training loop with K-Fold Cross-Validation
kf = KFold(n_splits=5, shuffle=True, random_state=42)
train_losses_per_epoch = np.zeros(n_epochs)
val_losses_per_epoch = np.zeros(n_epochs)
kf_count = 0
for train_idx, val_idx in kf.split(X_train):
X_kf_train, X_kf_val = X_train[train_idx], X_train[val_idx]
y_kf_train, y_kf_val = y_train[train_idx], y_train[val_idx]
for epoch in range(n_epochs):
model.train()
optimizer.zero_grad()
y_pred_train = model(X_kf_train)
train_loss = criterion(y_pred_train, y_kf_train)
train_loss.backward()
optimizer.step()
model.eval()
y_pred_val = model(X_kf_val)
val_loss = criterion(y_pred_val, y_kf_val)
train_losses_per_epoch[epoch] += train_loss.item()
val_losses_per_epoch[epoch] += val_loss.item()
kf_count += 1
# Average losses over K-Folds
train_losses_per_epoch /= kf_count
val_losses_per_epoch /= kf_count
# Evaluate on test data
model.eval()
y_pred_test = model(X_test)
test_loss = criterion(y_pred_test, y_test).item()
test_r2 = r2_score(y_test.detach().numpy(), y_pred_test.detach().numpy())
return train_losses_per_epoch, val_losses_per_epoch, test_loss, test_r2
# Run training and evaluation with TorchDynamo optimization
train_losses, val_losses, test_loss, test_r2 = train_and_evaluate(model, criterion, optimizer, X_train, y_train, X_test, y_test, n_epochs)
# Print metrics
print(f"Test MSE: {test_loss:.4f}")
print(f"Test R^2: {test_r2:.4f}")
# Plot results
epochs = list(range(1, n_epochs + 1))
plt.plot(epochs, train_losses, label='Train Loss')
plt.plot(epochs, val_losses, label='Validation Loss')
plt.xlabel('Epochs')
plt.ylabel('Loss')
plt.legend()
plt.title('Training and Validation Loss')
plt.show()

我们使用PyTorch定义了一个具有两个隐藏层的简单神经网络。模型使用K-Fold交叉验证来确保稳健的性能。TorchDynamo用于优化训练循环。在单独的测试集上对模型进行评估，并计算MSE和R²等指标。

最后得到的训练和验证损失如下

我们在代码中my_compiler打印了TorchDynamo相关的内容，我们来看看里面到底是什么：

my_compiler() called with FX graph:
opcode name target args kwargs
------------- ---------------------- ----------------------------------- ------------------------------------------------- --------
call_function train_losses_per_epoch

> (100,) {}
call_function val_losses_per_epoch

> (100,) {}
output output output ((train_losses_per_epoch, val_losses_per_epoch),) {}
my_compiler() called with FX graph:
opcode name target args kwargs
------------- ------------- ------------------------------------------------------- ---------------- --------
placeholder l_x_ L_x_ () {}
call_module l__self___fc1 L__self___fc1 (l_x_,) {}
call_function x

(l__self___fc1,) {}
call_module l__self___fc2 L__self___fc2 (x,) {}
call_function x_1

(l__self___fc2,) {}
call_module x_2 L__self___fc3 (x_1,) {}
output output output ((x_2,),) {}
my_compiler() called with FX graph:
opcode name target args kwargs
------------- ----------------------- -------------------------------------------------------- ------------------------------------------ --------------------------------
placeholder grad L_self_param_groups_0_params_0_grad () {}
placeholder grad_1 L_self_param_groups_0_params_1_grad () {}
placeholder grad_2 L_self_param_groups_0_params_2_grad () {}
placeholder grad_3 L_self_param_groups_0_params_3_grad () {}
placeholder grad_4 L_self_param_groups_0_params_4_grad () {}
placeholder grad_5 L_self_param_groups_0_params_5_grad () {}
get_attr param self___param_groups_0__params___0 () {}
get_attr param_1 self___param_groups_0__params___1 () {}
get_attr param_2 self___param_groups_0__params___2 () {}
get_attr param_3 self___param_groups_0__params___3 () {}

FX图的输出表明了模型的结构和操作是如何组织的:

输入0和L_x_是表示输入数据的占位符。

模型通过全连接层L__self___fc1 , L__self___fc2,L__self___fc3传递输入，这是神经网络的三层。

在前两层之后应用ReLU激活函数。

在第三层完全连接后产生最终输出。

总结

对于研究人员和工程师来说，训练大型和复杂的模型可能很耗时。TorchDynamo通过优化计算图和加速执行来减少这种训练时间，允许在更短的时间内进行更多的迭代和实验。在需要实时处理的应用程序中，如视频流或交互式人工智能系统，延迟是至关重要的。TorchDynamo在运行时优化和编译代码的能力确保了这些系统可以平稳运行并快速响应新数据。

TorchDynamo在支持多个后端和硬件架构方面的灵活性使其非常适合在各种环境中部署。无论是在高性能gpu或边缘设备上运行，TorchDynamo适应提供最佳性能。

https://avoid.overfit.cn/post/5fa68ddfeb024f569da6e09ae06c4ee4