用於部署的 TorchScript¶

創建於：2020 年 5 月 4 日 | 最後更新：2024 年 12 月 2 日 | 最後驗證：2024 年 11 月 5 日

警告

TorchScript 不再積極開發。

在本食譜中，您將學習

什麼是 TorchScript
如何以 TorchScript 格式導出您訓練好的模型
如何在 C++ 中加載您的 TorchScript 模型並進行推論

需求¶

PyTorch 1.5
TorchVision 0.6.0
libtorch 1.5
C++ 編譯器

在 pytorch.org 上可以找到安裝這三個 PyTorch 組件的說明。 C++ 編譯器將取決於您的平台。

什麼是 TorchScript？¶

TorchScript 是 PyTorch 模型（nn.Module 的子類）的中間表示形式，然後可以在像 C++ 這樣的效能環境中運行。它是 Python 的高效能子集，旨在被 PyTorch JIT 編譯器 使用，該編譯器對您模型的計算執行運行時優化。 TorchScript 是使用 PyTorch 模型進行縮放推論的推薦模型格式。有關更多信息，請參閱 PyTorch TorchScript 簡介教學、在 C++ 中加載 TorchScript 模型教學和完整的 TorchScript 文檔，所有這些都可以在 pytorch.org 上找到。

如何導出您的模型¶

作為一個例子，讓我們採用一個預訓練的視覺模型。 TorchVision 中的所有預訓練模型都與 TorchScript 兼容。

在腳本中或從 REPL 運行以下 Python 3 代碼

import torch
import torch.nn.functional as F
import torchvision.models as models

r18 = models.resnet18(pretrained=True)       # We now have an instance of the pretrained model
r18_scripted = torch.jit.script(r18)         # *** This is the TorchScript export
dummy_input = torch.rand(1, 3, 224, 224)     # We should run a quick test

讓我們對兩個模型的等效性進行健全性檢查

unscripted_output = r18(dummy_input)         # Get the unscripted model's prediction...
scripted_output = r18_scripted(dummy_input)  # ...and do the same for the scripted version

unscripted_top5 = F.softmax(unscripted_output, dim=1).topk(5).indices
scripted_top5 = F.softmax(scripted_output, dim=1).topk(5).indices

print('Python model top 5 results:\n  {}'.format(unscripted_top5))
print('TorchScript model top 5 results:\n  {}'.format(scripted_top5))

您應該看到這兩個版本的模型都給出相同的結果

Python model top 5 results:
  tensor([[463, 600, 731, 899, 898]])
TorchScript model top 5 results:
  tensor([[463, 600, 731, 899, 898]])

確認檢查後，繼續保存模型

r18_scripted.save('r18_scripted.pt')

在 C++ 中加載 TorchScript 模型¶

創建以下 C++ 文件並將其命名為 ts-infer.cpp

#include <torch/script.h>
#include <torch/nn/functional/activation.h>


int main(int argc, const char* argv[]) {
    if (argc != 2) {
        std::cerr << "usage: ts-infer <path-to-exported-model>\n";
        return -1;
    }

    std::cout << "Loading model...\n";

    // deserialize ScriptModule
    torch::jit::script::Module module;
    try {
        module = torch::jit::load(argv[1]);
    } catch (const c10::Error& e) {
        std::cerr << "Error loading model\n";
        std::cerr << e.msg_without_backtrace();
        return -1;
    }

    std::cout << "Model loaded successfully\n";

    torch::NoGradGuard no_grad; // ensures that autograd is off
    module.eval(); // turn off dropout and other training-time layers/functions

    // create an input "image"
    std::vector<torch::jit::IValue> inputs;
    inputs.push_back(torch::rand({1, 3, 224, 224}));

    // execute model and package output as tensor
    at::Tensor output = module.forward(inputs).toTensor();

    namespace F = torch::nn::functional;
    at::Tensor output_sm = F::softmax(output, F::SoftmaxFuncOptions(1));
    std::tuple<at::Tensor, at::Tensor> top5_tensor = output_sm.topk(5);
    at::Tensor top5 = std::get<1>(top5_tensor);

    std::cout << top5[0] << "\n";

    std::cout << "\nDONE\n";
    return 0;
}

這個程式

加載您在命令行上指定的模型
創建一個虛擬的“圖像”輸入張量
對輸入執行推論

另外，請注意此代碼中沒有對 TorchVision 的依賴。您保存的 TorchScript 模型具有您的學習權重和您的計算圖 - 不需要其他任何東西。

構建和運行您的 C++ 推論引擎¶

創建以下 CMakeLists.txt 文件

cmake_minimum_required(VERSION 3.0 FATAL_ERROR)
project(custom_ops)

find_package(Torch REQUIRED)

add_executable(ts-infer ts-infer.cpp)
target_link_libraries(ts-infer "${TORCH_LIBRARIES}")
set_property(TARGET ts-infer PROPERTY CXX_STANDARD 11)

製作程式

cmake -DCMAKE_PREFIX_PATH=<path to your libtorch installation>
make

現在，我們可以在 C++ 中運行推論，並驗證我們是否獲得了結果

$ ./ts-infer r18_scripted.pt
Loading model...
Model loaded successfully
 418
 845
 111
 892
 644
[ CPULongType{5} ]

DONE

重要資源¶

pytorch.org，獲取安裝說明，以及更多文檔和教程。
TorchScript 簡介教學，用於更深入地初步闡述 TorchScript
完整的 TorchScript 文檔，用於完整的 TorchScript 語言和 API 參考