torch.export 流程、常見挑戰及解決方案演示¶
作者： Ankith Gunapal, Jordi Ramon, Marcos Carranza
在`torch.export` 入門教程中，我們學習瞭如何使用torch.export。本教程在前一個教程的基礎上進行擴充套件，探討了使用程式碼匯出常用模型的過程，並解決了使用torch.export時可能出現的一些常見挑戰。
在本教程中，您將學習如何匯出適用於以下用例的模型：
影片分類器（MViT）
自動語音識別（OpenAI Whisper-Tiny）
影像字幕生成（BLIP）
可提示式影像分割（SAM2）
選擇這四種模型是為了演示 torch.export 的獨特功能，以及實現中遇到的一些實際考量和問題。
先決條件¶
PyTorch 2.4 或更高版本
對
torch.export和 PyTorch 即時推理有基本瞭解。
`torch.export` 的關鍵要求：無圖斷裂¶
torch.compile 透過使用 JIT 將 PyTorch 程式碼編譯成最佳化的核心來加速 PyTorch 程式碼。它使用 TorchDynamo 最佳化給定的模型，並建立最佳化的圖，然後使用 API 中指定的後端將其轉換為硬體特定的程式碼（lowered）。當 TorchDynamo 遇到不支援的 Python 特性時，它會中斷計算圖，讓預設的 Python 直譯器處理不支援的程式碼，然後恢復圖捕獲。計算圖中的這種中斷稱為圖斷裂。
`torch.export` 和 torch.compile 的關鍵區別之一在於 torch.export 不支援圖斷裂，這意味著您要匯出的整個模型或部分模型需要是一個單一的圖。這是因為處理圖斷裂涉及使用預設 Python 評估來解釋不受支援的操作，這與 torch.export 的設計目的不相容。您可以在此連結中閱讀有關不同 PyTorch 框架之間差異的詳細資訊
您可以使用以下命令來識別程式中的圖斷裂：
TORCH_LOGS="graph_breaks" python <file_name>.py
您需要修改程式以消除圖斷裂。解決後，您就可以匯出模型了。PyTorch 每晚都會在流行的 HuggingFace 和 TIMM 模型上執行 torch.compile 的基準測試。這些模型中的大多數都沒有圖斷裂。
本程式碼示例集中的模型沒有圖斷裂，但在使用 torch.export 時失敗。
影片分類¶
MViT 是一類基於 MultiScale Vision Transformers 的模型。該模型已使用 Kinetics-400 資料集訓練用於影片分類。該模型結合相關資料集可用於遊戲場景中的動作識別。
下面的程式碼透過使用 batch_size=2 進行跟蹤來匯出 MViT，然後檢查匯出的程式 (ExportedProgram) 是否可以使用 batch_size=4 執行。
import numpy as np
import torch
from torchvision.models.video import MViT_V1_B_Weights, mvit_v1_b
import traceback as tb
model = mvit_v1_b(weights=MViT_V1_B_Weights.DEFAULT)
# Create a batch of 2 videos, each with 16 frames of shape 224x224x3.
input_frames = torch.randn(2,16, 224, 224, 3)
# Transpose to get [1, 3, num_clips, height, width].
input_frames = np.transpose(input_frames, (0, 4, 1, 2, 3))
# Export the model.
exported_program = torch.export.export(
model,
(input_frames,),
)
# Create a batch of 4 videos, each with 16 frames of shape 224x224x3.
input_frames = torch.randn(4,16, 224, 224, 3)
input_frames = np.transpose(input_frames, (0, 4, 1, 2, 3))
try:
exported_program.module()(input_frames)
except Exception:
tb.print_exc()
錯誤：靜態批次大小¶
raise RuntimeError(
RuntimeError: Expected input at *args[0].shape[0] to be equal to 2, but got 4
預設情況下，匯出流程會在跟蹤程式時假定所有輸入形狀都是靜態的，因此如果您使用與跟蹤時不同的輸入形狀執行程式，將會遇到錯誤。
解決方案¶
為解決該錯誤，我們將輸入的第一維（batch_size）指定為動態的，並指定 batch_size 的預期範圍。在下面顯示的更正後的示例中，我們指定預期的 batch_size 範圍為 1 到 16。需要注意的一個細節是 min=2 並不是一個 bug，並在The 0/1 Specialization Problem中進行了解釋。關於 torch.export 動態形狀的詳細說明可以在匯出教程中找到。下面顯示的程式碼演示瞭如何匯出具有動態批次大小的 mViT：
import numpy as np
import torch
from torchvision.models.video import MViT_V1_B_Weights, mvit_v1_b
import traceback as tb
model = mvit_v1_b(weights=MViT_V1_B_Weights.DEFAULT)
# Create a batch of 2 videos, each with 16 frames of shape 224x224x3.
input_frames = torch.randn(2,16, 224, 224, 3)
# Transpose to get [1, 3, num_clips, height, width].
input_frames = np.transpose(input_frames, (0, 4, 1, 2, 3))
# Export the model.
batch_dim = torch.export.Dim("batch", min=2, max=16)
exported_program = torch.export.export(
model,
(input_frames,),
# Specify the first dimension of the input x as dynamic
dynamic_shapes={"x": {0: batch_dim}},
)
# Create a batch of 4 videos, each with 16 frames of shape 224x224x3.
input_frames = torch.randn(4,16, 224, 224, 3)
input_frames = np.transpose(input_frames, (0, 4, 1, 2, 3))
try:
exported_program.module()(input_frames)
except Exception:
tb.print_exc()
自動語音識別¶
自動語音識別 (ASR) 是指使用機器學習將口語轉錄成文字。Whisper 是 OpenAI 基於 Transformer 的編碼器-解碼器模型，該模型使用 68 萬小時的標註資料進行 ASR 和語音翻譯的訓練。下面的程式碼嘗試匯出用於 ASR 的 whisper-tiny 模型。
import torch
from transformers import WhisperProcessor, WhisperForConditionalGeneration
from datasets import load_dataset
# load model
model = WhisperForConditionalGeneration.from_pretrained("openai/whisper-tiny")
# dummy inputs for exporting the model
input_features = torch.randn(1,80, 3000)
attention_mask = torch.ones(1, 3000)
decoder_input_ids = torch.tensor([[1, 1, 1 , 1]]) * model.config.decoder_start_token_id
model.eval()
exported_program: torch.export.ExportedProgram= torch.export.export(model, args=(input_features, attention_mask, decoder_input_ids,))
錯誤：使用 TorchDynamo 進行嚴格跟蹤¶
torch._dynamo.exc.InternalTorchDynamoError: AttributeError: 'DynamicCache' object has no attribute 'key_cache'
預設情況下，torch.export 使用TorchDynamo（一種位元組碼分析引擎）跟蹤您的程式碼，它對您的程式碼進行符號分析並構建圖。這種分析提供了更強的安全性保證，但並非所有 Python 程式碼都受支援。當我們使用預設的嚴格模式匯出 whisper-tiny 模型時，通常會在 Dynamo 中因不支援的特性而返回錯誤。要了解為何在 Dynamo 中出現此錯誤，您可以參考此GitHub issue。
解決方案¶
為解決上述錯誤，torch.export 支援 non_strict 模式，在該模式下，程式使用 Python 直譯器進行跟蹤，其工作方式類似於 PyTorch 即時執行。唯一的區別是所有 Tensor 物件都將被 ProxyTensors 替換，ProxyTensors 會將它們的所有操作記錄到一個圖中。透過使用 strict=False，我們能夠匯出程式。
import torch
from transformers import WhisperProcessor, WhisperForConditionalGeneration
from datasets import load_dataset
# load model
model = WhisperForConditionalGeneration.from_pretrained("openai/whisper-tiny")
# dummy inputs for exporting the model
input_features = torch.randn(1,80, 3000)
attention_mask = torch.ones(1, 3000)
decoder_input_ids = torch.tensor([[1, 1, 1 , 1]]) * model.config.decoder_start_token_id
model.eval()
exported_program: torch.export.ExportedProgram= torch.export.export(model, args=(input_features, attention_mask, decoder_input_ids,), strict=False)
影像字幕生成¶
影像字幕生成 是將影像內容用文字描述的任務。在遊戲場景中，影像字幕生成可用於透過動態生成場景中各種遊戲物件的文字描述來增強遊戲體驗，從而為玩家提供額外細節。BLIP 是 SalesForce Research 釋出的一個流行的影像字幕生成模型。下面的程式碼嘗試使用 batch_size=1 匯出 BLIP。
import torch
from models.blip import blip_decoder
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
image_size = 384
image = torch.randn(1, 3,384,384).to(device)
caption_input = ""
model_url = 'https://storage.googleapis.com/sfr-vision-language-research/BLIP/models/model_base_capfilt_large.pth'
model = blip_decoder(pretrained=model_url, image_size=image_size, vit='base')
model.eval()
model = model.to(device)
exported_program: torch.export.ExportedProgram= torch.export.export(model, args=(image,caption_input,), strict=False)
錯誤：無法修改具有凍結儲存的張量¶
在匯出模型時，可能會失敗，因為模型實現可能包含 torch.export 尚未支援的某些 Python 操作。其中一些失敗可能有變通方法。BLIP 就是一個例子，原始模型會出錯，可以透過對程式碼進行少量修改來解決。torch.export 在 ExportDB 中列出了支援和不支援操作的常見情況，並展示瞭如何修改程式碼以使其相容匯出。
File "/BLIP/models/blip.py", line 112, in forward
text.input_ids[:,0] = self.tokenizer.bos_token_id
File "/anaconda3/envs/export/lib/python3.10/site-packages/torch/_subclasses/functional_tensor.py", line 545, in __torch_dispatch__
outs_unwrapped = func._op_dk(
RuntimeError: cannot mutate tensors with frozen storage
可提示式影像分割¶
影像分割是一種計算機視覺技術，根據畫素的特徵將其分為不同的畫素組或區域（segments）。Segment Anything Model (SAM)) 引入了可提示式影像分割，它根據指示所需物件的提示來預測物件掩碼。SAM 2 是第一個用於跨影像和影片分割物件的統一模型。SAM2ImagePredictor 類為模型提供了一個簡單的介面，用於對模型進行提示。該模型可以同時接受點和框提示作為輸入，也可以接受前一次預測產生的掩碼作為輸入。由於 SAM2 在目標跟蹤方面提供了強大的零樣本效能，因此可用於跟蹤場景中的遊戲物件。
在 SAM2ImagePredictor 的 predict 方法中的張量操作在 _predict 方法中進行。因此，我們嘗試像這樣匯出：
ep = torch.export.export(
self._predict,
args=(unnorm_coords, labels, unnorm_box, mask_input, multimask_output),
kwargs={"return_logits": return_logits},
strict=False,
)
錯誤：模型不是 torch.nn.Module 型別¶
torch.export 要求模組型別為 torch.nn.Module。然而，我們嘗試匯出的模組是一個類方法。因此會出錯。
Traceback (most recent call last):
File "/sam2/image_predict.py", line 20, in <module>
masks, scores, _ = predictor.predict(
File "/sam2/sam2/sam2_image_predictor.py", line 312, in predict
ep = torch.export.export(
File "python3.10/site-packages/torch/export/__init__.py", line 359, in export
raise ValueError(
ValueError: Expected `mod` to be an instance of `torch.nn.Module`, got <class 'method'>.
解決方案¶
我們編寫一個輔助類，該類繼承自 torch.nn.Module，並在該類的 forward 方法中呼叫 _predict 方法。完整的程式碼可以在這裡找到。
class ExportHelper(torch.nn.Module):
def __init__(self):
super().__init__()
def forward(_, *args, **kwargs):
return self._predict(*args, **kwargs)
model_to_export = ExportHelper()
ep = torch.export.export(
model_to_export,
args=(unnorm_coords, labels, unnorm_box, mask_input, multimask_output),
kwargs={"return_logits": return_logits},
strict=False,
)
結論¶
在本教程中，我們學習瞭如何透過正確的配置和簡單的程式碼修改來解決挑戰，從而使用 torch.export 匯出適用於流行用例的模型。成功匯出模型後，對於伺服器，您可以使用 AOTInductor 將 ExportedProgram 轉換為硬體可執行程式碼；對於邊緣裝置，則可以使用 ExecuTorch。要了解有關 AOTInductor (AOTI) 的更多資訊，請參考 AOTI 教程。要了解有關 ExecuTorch 的更多資訊，請參考 ExecuTorch 教程。
