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 develop

在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("numpy.**")
    e.extern("sys")
    e.extern("PIL.*")
    e.save_pickle("model", "model.pkl", model)

请注意，由于“numpy”、“sys”和“PIL”被标记为“extern”，torch.package将在加载此包的系统上查找这些依赖项。它们不会与模型一起打包。

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

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

设置一个环境变量（例如 $PATH_TO_EXTERN_PYTHON_PACKAGES），以告知解释器在哪里可以找到外部 Python 依赖项。在下面的示例中，提供了一个 conda 环境的 site-packages 路径。

export PATH_TO_EXTERN_PYTHON_PACKAGES= \
    "~/anaconda/envs/deploy-example-env/lib/python3.8/site-packages"

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

#include <torch/csrc/deploy/deploy.h>
#include <torch/csrc/deploy/path_environment.h>
#include <torch/script.h>
#include <torch/torch.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.
    std::shared_ptr<torch::deploy::Environment> env =
        std::make_shared<torch::deploy::PathEnvironment>(
            std::getenv("PATH_TO_EXTERN_PYTHON_PACKAGES")
        );
    torch::deploy::InterpreterManager manager(4, env);

    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";
        std::cerr << e.msg();
        return -1;
    }

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

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

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

PathEnvironment 使你能够指定系统上 Python 包的路径，这些包是外部的，但对你的模型是必要的。

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

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

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

构建和运行应用程序

在您的系统上找到 libtorch_deployinterpreter.o。这应该是在从源代码构建 PyTorch 时一并生成的。在同一个 PyTorch 目录中，找到部署源文件。将这些位置设置为构建的环境变量。下面显示了这些文件在系统上的一个示例位置。

export DEPLOY_INTERPRETER_PATH="/pytorch/build/torch/csrc/deploy/"
export DEPLOY_SRC_PATH="/pytorch/torch/csrc/deploy/"

由于 torch::deploy 处于积极开发中，这些手动步骤将很快被移除。

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

cmake_minimum_required(VERSION 3.19 FATAL_ERROR)
project(deploy_tutorial)

find_package(fmt REQUIRED)
find_package(Torch REQUIRED)

add_library(torch_deploy_internal STATIC
    ${DEPLOY_INTERPRETER_PATH}/libtorch_deployinterpreter.o
    ${DEPLOY_DIR}/deploy.cpp
    ${DEPLOY_DIR}/loader.cpp
    ${DEPLOY_DIR}/path_environment.cpp
    ${DEPLOY_DIR}/elf_file.cpp)

# for python builtins
target_link_libraries(torch_deploy_internal PRIVATE
    crypt pthread dl util m z ffi lzma readline nsl ncursesw panelw)
target_link_libraries(torch_deploy_internal PUBLIC
    shm torch fmt::fmt-header-only)
caffe2_interface_library(torch_deploy_internal torch_deploy)

add_executable(example-app example.cpp)
target_link_libraries(example-app PUBLIC
    "-Wl,--no-as-needed -rdynamic" dl torch_deploy "${TORCH_LIBRARIES}")

目前，需要将 torch::deploy 构建为静态库。 为了正确链接到静态库，使用实用工具 caffe2_interface_library 来正确设置和取消设置 --whole-archive 标志。

此外，当链接到可执行文件时，需要使用 -rdynamic 标志 以确保符号被导出到动态表中，从而使它们可以被 部署解释器（动态加载）访问。

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

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

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

mkdir build
cd build
# Point CMake at the built version of PyTorch we just installed.
cmake -DCMAKE_PREFIX_PATH="$(python -c 'import torch.utils; print(torch.utils.cmake_prefix_path)')" .. \
    -DDEPLOY_INTERPRETER_PATH="$DEPLOY_INTERPRETER_PATH" \
    -DDEPLOY_DIR="$DEPLOY_DIR"
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()，并且前向方法将在多个独立的解释器上执行，而不会受到全局解释器锁的限制。