用于部署的TorchScript¶

Created On: May 04, 2020 | Last Updated: Dec 02, 2024 | Last Verified: Nov 05, 2024

警告

TorchScript 不再被积极开发。

在此教程中，您将学习：

什么是TorchScript
如何以TorchScript格式导出您的训练模型
如何在C++中加载TorchScript模型并进行推理

要求¶

PyTorch 1.5
TorchVision 0.6.0
libtorch 1.5
C++编译器

安装三个PyTorch组件的说明可在 pytorch.org 上找到。C++编译器将取决于您的平台。

什么是TorchScript?¶

TorchScript 是PyTorch模型（“nn.Module”的子类）的中间表示，可以在像C++这样的高性能环境中运行。它是一种高性能的Python子集，旨在被 PyTorch JIT编译器 使用，该编译器对模型的计算进行运行时优化。TorchScript是使用PyTorch模型进行扩展推理的推荐模型格式。有关更多信息，请参见PyTorch TorchScript介绍教程、在C++中加载TorchScript模型教程和完整的 TorchScript文档，它们均可在 pytorch.org 上找到。

如何导出您的模型¶

作为示例，我们采用一个预训练视觉模型。TorchVision中的所有预训练模型都兼容TorchScript。

运行以下Python 3代码，无论是在脚本中还是在REPL中:

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;
}

该程序:

加载您在命令行中指定的模型
创建一个dummy“图像”输入张量
对输入进行推理

此外，请注意，此代码中没有对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参考