Giter Club home page Giter Club logo

pytorch2tensorrt's Introduction

  • Pytorch2TensorRT

将Pytorch模型部署到TensorRT的一个简单用法,技术路线为“pytorch model-->onnx file-->TensorRT engine”。

当前仅针对ONNX和TensorRT支持OP可进行转换,如有不支持的OP需编写插件。

News:

  • 2020.12.10: 更新trt_convertor.py脚本,使之适用于TRT7;

软件环境:

TensorRT7.0.0.11
Pytorch1.2
PIL6.2.1
numpy1.17.4

Linux_x86_64
CUDA10.0
CUDNN7.5.0

当前支持:

  • TensorRT FP32
  • TensorRT FP16
  • TensorRT INT8

使用方法:

  1. 从Pytorch模型到ONNX:修改并使用pytoch_to_onnx.py脚本转ONNX,或者独自进行转换;
  2. 利用自行提供的或根据上一步转换好的ONNX文件,进行TensorRT转换:Python main.py,并指定必要的参数;
  3. 使用do_inference.py进行推理验证。

使用示例:

ONNX file to FP16 engine: python main.py --batch_size 32 --mode fp16 --onnx_file_path my_files/centernet.onnx --engine_file_path my_files/test_fp16.engine

推理:python do_inference.py --engine_file_path my_files/test.engine --img_path test_img.jpg --batch_size 1

使用说明:

Pytorch模型转ONNX:

  • 参考脚本pytoch_to_onnx.py,需按照自己的需要定义模型与输入样例,然后转换。

将ONNX转换为INT8的TensorRT引擎,需要:

  1. 准备一个校准集,用于在转换过程中寻找使得转换后的激活值分布与原来的FP32类型的激活值分布差异最小的阈值;
  2. 并写一个校准器类,该类需继承trt.IInt8EntropyCalibrator2父类,并重写get_batch_size, get_batch, read_calibration_cache, write_calibration_cache这几个方法。具体做法参考脚本myCalibrator.py.
  3. 使用时,需额外指定cache_file,该参数是校准集cache文件的路径,会在校准过程中生成,方便下一次校准时快速提取。

参考:

https://github.com/GuanLianzheng/pytorch_to_TensorRT5.git

官方示例:path_to_tensorrt/TensorRT-5.1.5.0/samples/python/int8_caffe_mnist

pytorch2tensorrt's People

Contributors

ailiwensong avatar

Recommend Projects

  • React photo React

    A declarative, efficient, and flexible JavaScript library for building user interfaces.

  • Vue.js photo Vue.js

    🖖 Vue.js is a progressive, incrementally-adoptable JavaScript framework for building UI on the web.

  • Typescript photo Typescript

    TypeScript is a superset of JavaScript that compiles to clean JavaScript output.

  • TensorFlow photo TensorFlow

    An Open Source Machine Learning Framework for Everyone

  • Django photo Django

    The Web framework for perfectionists with deadlines.

  • D3 photo D3

    Bring data to life with SVG, Canvas and HTML. 📊📈🎉

Recommend Topics

  • javascript

    JavaScript (JS) is a lightweight interpreted programming language with first-class functions.

  • web

    Some thing interesting about web. New door for the world.

  • server

    A server is a program made to process requests and deliver data to clients.

  • Machine learning

    Machine learning is a way of modeling and interpreting data that allows a piece of software to respond intelligently.

  • Game

    Some thing interesting about game, make everyone happy.

Recommend Org

  • Facebook photo Facebook

    We are working to build community through open source technology. NB: members must have two-factor auth.

  • Microsoft photo Microsoft

    Open source projects and samples from Microsoft.

  • Google photo Google

    Google ❤️ Open Source for everyone.

  • D3 photo D3

    Data-Driven Documents codes.