Skip to content

Latest commit

 

History

History
44 lines (32 loc) · 763 Bytes

README.md

File metadata and controls

44 lines (32 loc) · 763 Bytes

Pytorch-Model-to-TensorRT

Convert a Pytorch model to TensorRT engin (optional: int8 engin)

Required:

Python packages: in requirements.txt

External:

  • CUDA == 9.0
  • CUDNN == 7.3.1
  • TensorRT == 4.0.2.6 link

Data And Checkpoint Prepare

You should prepare data and model weights in like this way:

project  
│
└───data
│   │   test.txt
│   │   fake-test.txt
│   │   images
│   
└───checkpoint
    │   model_best.pth

The list of test file should follow this format:

file0,label0
file1,label1
... ...

Or you can download the example from BaiduYun

Convert step:

python main.py