Skip to content

Latest commit

 

History

History
 
 

psenet

PSENet

preprocessing + inference + postprocessing = 30ms with fp32 on Tesla P40. The original Tensorflow implementation is tensorflow_PSENet. A TensorRT Python api implementation is TensorRT-Python-PSENet.

Key Features

  • Generating .wts from Tensorflow.
  • Dynamic batch and dynamic shape input.
  • Object-Oriented Programming.
  • Practice with C++ 11.

How to Run

    1. generate .wts

    Download pretrained model from https://github.com/liuheng92/tensorflow_PSENet and put model.ckpt.* to model dir. Add a file model/checkpoint with content

    model_checkpoint_path: "model.ckpt"
    all_model_checkpoint_paths: "model.ckpt"
    

    Then run

    python gen_tf_wts.py
    

    which will gengerate a psenet.wts.

    1. cmake and make
    mkdir build
    cd build
    cmake ..
    make
    
    1. build engine and run detection
    cp ../psenet.wts ./
    cp ../test.jpg ./
    ./psenet -s  // serialize model to plan file
    ./psenet -d  // deserialize plan file and run inference
    

Known Issues

None

Todo

  • use ExponentialMovingAverage weight.