专栏名称: AI算法与图像处理
考研逆袭985,非科班跨行AI,目前从事计算机视觉的工业和商业相关应用的工作。分享最新最前沿的科技,共同分享宝贵的资源资料,这里有机器学习,计算机视觉,Python等技术实战分享,也有考研,转行IT经验交流心得
目录
相关文章推荐
宝鸡市场监管  ·  这样洗草莓可能农残更高!很多人都做错了…… ·  19 小时前  
江苏药品监管  ·  不是错觉!过了这个年纪,真的更容易长胖 ·  2 天前  
51好读  ›  专栏  ›  AI算法与图像处理

Pytorch2.x时代,关于C++部署的讨论

AI算法与图像处理  · 公众号  ·  · 2024-11-06 21:32

正文

Pytorch模型的 高性能部署 一直是大家讨论的问题,有两点比较重要:
  • 高度优化的算子
  • 可以高效率运行计算图的架构和runtime

高度优化的算子不用多说,TensorRT为什么那么快,因为engine在构建的时候,在每个平台(A10、A100、T4等)上搜索到了最优最快的kernel(实现了一些op)。高效率运行计算图也是很关键的一点,TensorRT构建好engine后,需要libnvinfer.so来驱动,其中实现了什么,在使用过程中很容易猜到:

  • 序列化和反序列化,也就是所谓的生成engine,读取engine
  • 推理engine、多stream运行计算图,管理engine所需要的一些环境,比如显存和中间变量等

为了达到极致的性能,TensorRT的整个运行时都是在C++环境中,虽然提供了Python-API,但实际调用执行的操作都是在C++中,Python只提供包了一层的作用,算子和执行整个计算图的地方都是C++。

c++ api vs python api

python有快速开发以及验证的优点,但是相比C++来说 速度较慢 而且比较 费内存 ,一般高性能场景都是使用C++去部署,尽量避免使用python环境。

TORCH 1.x时期的C++部署

torch1.x的实际场景中,一般是搭配使用libtorch + torchscript,这俩在很多生产环境中已经验证过了。

libtorch可以使用C++ API去完成和python中使用pytorch-op实现一样的功能,比如:

#include 

at::Tensor a = at::ones({22}, at::kInt);
at::Tensor b = at::randn({22});
auto c = a + b.to(at::kInt);

转化为Pytorch就是:

import torch

a = torch.ones((22), dtype=torch.int32)
b = torch.randn((22))
c = a + b.to(torch.int32)

而torchscript则用于 trace或者script我们的模型到C++环境中 部署,速度方面变化不大,主要是通过torchscript导出的模型可以在C++环境中加载并运行,不需要依赖python了,可以减少一些python的over head:

#include  // One-stop header.

#include 
#include 

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


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

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

关于torchscript的解读有不少,这里不赘述了,感兴趣的可以参阅:

  • https://zhuanlan.zhihu.com/p/486914187
  • https://zhuanlan.zhihu.com/p/489090393
  • https://zhuanlan.zhihu.com/p/363319763
  • https://zhuanlan.zhihu.com/p/652193676
  • https://zhuanlan.zhihu.com/p/410507557

TORCH 2.x的C++部署

torch2.0出来的时候,最主要的就是torch.compile的新API,可以直接优化模型。

torch.compile核心是dynamo, dynamo相比torch.jit.trace和torch.jit.script,是一个功能更强大的trace工具 [1] ,trace模型从而优化模型。dynamo出现后,我也很好奇torchscript是否会被废弃?

torchscript

目前看来torchscript还是会继续存在,只是freeze了,功能还会维护,bug还会修,但不会有新功能了。

基于torch.jit.trace的模型导出路径成为过去式了,那么新的基于pt2.0的C++导出方案是啥?

torch官方前一周发了一篇新blog,正式提到了cpp wrapper,核心就是 torch.export [2] + cpp wrapper [3]

  • PyTorch 2.1 Contains New Performance Features for AI Developers [4]

使用cpp wrapper去invoke the generated kernels and external kernels in TorchInductor,可以减少python的overhead,实际测试中,模型速度越快,python overhead占比越大,提升也就越大:

cpp wrapper benchmark

我们都知道torch2.0可以基于triton生成高性能的kernel,例如:

@torch.compile
def opt_foo(x, y):
    a = torch.sin(x)
    b = torch.cos(y)
    return a + b

for _ in range(100):
    opt_foo(torch.randn(10).cuda(), torch.randn(10).cuda())

定义好一个函数后,加上 @torch.compile 装饰器,执行几次即可得到优化后的模型,默认使用的优化器是TorchInductor,借助 depyf [5] ,我们可以看到优化好后生成的triton代码(GPU端):

import triton
import triton.language as tl
from torch._inductor.ir import ReductionHint
from torch._inductor.ir import TileHint
from torch._inductor.triton_heuristics import AutotuneHint, pointwise
from torch._inductor.utils import instance_descriptor
from torch._inductor import triton_helpers

@pointwise(
    size_hints=[16], 
    filename=__file__,
    triton_meta={'signature': {0'*fp32'1'*fp32'2'*fp32'3'i32'}, 'device'0'device_type''cuda''constants': {}, 'configs': [instance_descriptor(divisible_by_16=(012), equal_to_1=(), ids_of_folded_args=(), divisible_by_8=())]},
    inductor_meta={'autotune_hints': set(), 'kernel_name''triton_poi_fused_add_cos_sin_0''mutated_arg_names': []},
    min_elem_per_thread=0
)
@triton.jit
def triton_(in_ptr0, in_ptr1, out_ptr0, xnumel, XBLOCK : tl.constexpr):
    xnumel = 10
    xoffset = tl.program_id(0) * XBLOCK
    xindex = xoffset + tl.arange(0, XBLOCK)[:]
    xmask = xindex     x0 = xindex
    tmp0 = tl.load(in_ptr0 + (x0), xmask)
    tmp2 = tl.load(in_ptr1 + (x0), xmask)
    tmp1 = tl.sin(tmp0)
    tmp3 = tl.cos(tmp2)
    tmp4 = tmp1 + tmp3
    tl.store(out_ptr0 + (x0), tmp4, xmask)

这个triton代码可以直接调用,但是依赖python环境,如果想要切换到C++端,则修改下config:

import torch._inductor.config as config
config.cpp_wrapper = True

后重新执行几次,可以得到生成的cpp调用代码:

#include 
#include 
#include 
#include 
#include 
#define reinterpret_tensor torch::inductor::_reinterpret_tensor
#define alloc_from_pool torch::inductor::_alloc_from_pool
#include 

[[maybe_unused]] static int64_t align(int64_t nbytes) {
  return (nbytes + 64 - 1) & -64;
}
#include 

#include 
#include 

#define CUDA_DRIVER_CHECK(EXPR)                    \
do {                                               \
    CUresult code = EXPR;                          \
    const char *msg;                               \
    cuGetErrorString(code, &msg);                  \
    if (code != CUDA_SUCCESS) {                    \
        throw std::runtime_error(                  \
            std::string("CUDA driver error: ") +   \
            std::string(msg));                     \
    }                                              \
} while (0);


namespace {

struct Grid {
    Grid(uint32_t x, uint32_t y, uint32_t z)
      : grid_x(x), grid_y(y), grid_z(z) {}
    uint32_t grid_x;
    uint32_t grid_y;
    uint32_t grid_z;

    bool is_non_zero() {
        return grid_x > 0 && grid_y > 0 && grid_z > 0;
    }
};

}  // anonymous namespace

static inline CUfunction loadKernel(
        std::string






请到「今天看啥」查看全文