深度学习系列72:torch-tensorrt入门

1. 安装

坑非常多,清华源阿里源都不行。使用官网源下载,这里的121可以改成你需要的东西:

python -m pip install torch torch-tensorrt tensorrt --extra-index-url https://download.pytorch.org/whl/cu121

2. 原理

我们来看一个实例:这是一个用于支持 torchscript 到 TensorRT 转换的项目。上面的代码用于将 addmm 运算展开成数个算子,方便后续映射 TensorRT 算子。

void UnpackAddMM(std::shared_ptr<torch::jit::Graph>& graph) {
  // TensorRT implicitly adds a flatten layer in front of FC layers if necessary
  // 用于匹配的模式
  std::string addmm_pattern = R"IR(
    graph(%b, %x, %w, %beta, %alpha):
      %out: Tensor = aten::addmm(%b, %x, %w, %beta, %alpha)
      return (%out))IR";
  // 用于替换的模式
  std::string mm_add_pattern = R"IR(
    graph(%b, %x, %w, %beta, %alpha):
      %mm: Tensor = aten::matmul(%x, %w)
      %bias: Tensor = aten::mul(%b, %beta)
      %out: Tensor = aten::add(%bias, %mm, %alpha)
      return (%out))IR";

  // 创建子图重写器并注册匹配模式和替换模式
  torch::jit::SubgraphRewriter unpack_addmm;
  unpack_addmm.RegisterRewritePattern(addmm_pattern, mm_add_pattern);
  // 遍历graph,完成重写
  unpack_addmm.runOnGraph(graph);
  LOG_GRAPH("Post unpack addmm: " << *graph);
}

3. 简单例子

import torch
def origin_func(x):
    x = x**2
    x = x**3
    return x

x = torch.rand(1, 2, 3, 4)
jit_model = torch.jit.trace(origin_func, x)
print(jit_model.graph)

# 匹配用的子图定义,注意常量必须为[value=2]属性
pattern = """
    graph(%x):
        %const_2 = prim::Constant[value=2]()
        %out = aten::pow(%x, %const_2)
        return (%out)
"""
# 替换用的子图定义
replacement = """
    graph(%x):
        %out = aten::mul(%x, %x)
        return (%out)
"""
torch._C._jit_pass_custom_pattern_based_rewrite_graph(pattern, replacement,jit_model.graph)
print(jit_model.graph)
相关推荐
东胜物联10 分钟前
探寻5G工业网关市场,5G工业网关品牌解析
人工智能·嵌入式硬件·5g
皓74121 分钟前
服饰电商行业知识管理的创新实践与知识中台的重要性
大数据·人工智能·科技·数据分析·零售
wangyue438 分钟前
c# 深度模型入门
深度学习
川石课堂软件测试1 小时前
性能测试|docker容器下搭建JMeter+Grafana+Influxdb监控可视化平台
运维·javascript·深度学习·jmeter·docker·容器·grafana
985小水博一枚呀1 小时前
【深度学习滑坡制图|论文解读3】基于融合CNN-Transformer网络和深度迁移学习的遥感影像滑坡制图方法
人工智能·深度学习·神经网络·cnn·transformer
AltmanChan1 小时前
大语言模型安全威胁
人工智能·安全·语言模型
985小水博一枚呀1 小时前
【深度学习滑坡制图|论文解读2】基于融合CNN-Transformer网络和深度迁移学习的遥感影像滑坡制图方法
人工智能·深度学习·神经网络·cnn·transformer·迁移学习
数据与后端架构提升之路1 小时前
从神经元到神经网络:深度学习的进化之旅
人工智能·神经网络·学习
爱技术的小伙子1 小时前
【ChatGPT】如何通过逐步提示提高ChatGPT的细节描写
人工智能·chatgpt
深度学习实战训练营3 小时前
基于CNN-RNN的影像报告生成
人工智能·深度学习