torch::deploy

torch::deploy 是一个系统，允许你在 C++ 进程中运行多个嵌入式 Python 解释器，而无需共享全局解释器锁。有关 torch::deploy 内部工作原理的更多信息，请参阅相关 arXiv 论文。

警告

这是一个原型功能。仅支持 Linux x86，API 可能在未提前通知的情况下发生变化。

入门

安装 torch::deploy

torch::deploy 在我们的二进制发布版本中尚未默认构建，因此要获取启用了 torch::deploy 的 libtorch 副本，请按照 从源代码构建 PyTorch 的说明进行操作。

当运行 setup.py 时，你需要指定 USE_DEPLOY=1，例如:

export CMAKE_PREFIX_PATH=${CONDA_PREFIX:-"$(dirname $(which conda))/../"}
export USE_DEPLOY=1
python setup.py bdist_wheel
python -mpip install dist/*.whl

在Python中创建模型包

torch::deploy 可以加载和运行使用 torch.package 打包的 Python 模型。您可以在 torch.package 的 torch.package 文档 中了解更多。

目前，让我们创建一个简单的模型，可以在 torch::deploy 中加载并运行。

from torch.package import PackageExporter
import torchvision

# Instantiate some model
model = torchvision.models.resnet.resnet18()

# Package and export it.
with PackageExporter("my_package.pt") as e:
    e.intern("torchvision.**")
    e.extern("sys")
    e.save_pickle("model", "model.pkl", model)

现在，你的工作目录中应该有一个名为 my_package.pt 的文件。

注意

目前，torch::deploy 仅支持 Python 标准库和 torch 作为 extern 模块在 torch.package 中。未来我们计划 透明地支持您指向我们的任何 Conda 环境。

在C++中加载和运行模型

让我们创建一个最简的 C++ 程序来加载该模型。

#include <torch/deploy.h>
#include <torch/script.h>

#include <iostream>
#include <memory>

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

    // Start an interpreter manager governing 4 embedded interpreters.
    torch::deploy::InterpreterManager manager(4);

    try {
        // Load the model from the torch.package.
        torch::deploy::Package package = manager.loadPackage(argv[1]);
        torch::deploy::ReplicatedObj model = package.loadPickle("model", "model.pkl");
    } catch (const c10::Error& e) {
        std::cerr << "error loading the model\n";
        return -1;
    }

    std::cout << "ok\n";
}

这个小程序介绍了torch::deploy的许多核心概念。

An InterpreterManager 抽象了多个独立的 Python 解释器集合，允许你在运行代码时在它们之间进行负载均衡。

使用 InterpreterManager::loadPackage 方法，你可以从磁盘加载一个 torch.package 并使其对所有解释器可用。

Package::loadPickle 允许你从该包中检索特定的 Python 对象，比如我们之前保存的 ResNet 模型。

最后，模型本身是一个 ReplicatedObj。这是一个抽象的句柄，指向一个在多个解释器中复制的对象。当你与一个 ReplicatedObj 进行交互（例如，通过调用 forward），它会选择一个空闲的解释器来执行该交互。

构建和运行应用程序

假设上述C++程序存储在一个名为 example-app.cpp 的文件中，一个最小的 CMakeLists.txt 文件将如下所示：

cmake_minimum_required(VERSION 3.0 FATAL_ERROR)
project(deploy_tutorial)

find_package(Torch REQUIRED)

add_executable(example-app example-app.cpp)
target_link_libraries(example-app "${TORCH_LIBRARIES}")
set_property(TARGET example-app PROPERTY CXX_STANDARD 14)

最后一步是配置和构建项目。假设我们的代码目录布局如下：

example-app/
    CMakeLists.txt
    example-app.cpp

我们现在可以运行以下命令，从example-app/文件夹内构建应用程序:

mkdir build
cd build
# Point CMake at the built version of PyTorch we just installed.
SITE_PACKAGES="$(python -c 'from distutils.sysconfig import get_python_lib; print(get_python_lib())')"
cmake -DCMAKE_PREFIX_PATH="$SITE_PACKAGES/torch" ..
cmake --build . --config Release

现在我们可以运行我们的应用：

./example-app /path/to/my_package.pt

执行 forward 的 C++

一旦你将模型加载到 C++ 中，执行它就变得很容易：

// Create a vector of inputs.
std::vector<torch::jit::IValue> inputs;
inputs.push_back(torch::ones({1, 3, 224, 224}));

// Execute the model and turn its output into a tensor.
at::Tensor output = model(inputs).toTensor();
std::cout << output.slice(/*dim=*/1, /*start=*/0, /*end=*/5) << '\n';

值得注意的是，模型的前向函数是在Python中执行的，使用的是嵌入式的CPython解释器。请注意，该模型是一个ReplicatedObj，这意味着你可以从多个线程调用model()，并且前向方法将在多个独立的解释器上执行，而不会受到全局解释器锁的限制。