The Pytorch implementation is HRNet-Image-Classification. The implemented model is HRNet-W18-C-Small-v2

• 1. generate .wts

  Download code and model from HRNet-Image-Classification and config your environments.

  Put demo.py in the YOUR_ROOT_DIR\HRNet-Image-Classification\tools folder, set savewts in main() as True, and run, the .wts will be generated.

• 2. cmake and make

  mkdir build
  cd build
  cmake ..
  make
  sudo ./hrnet -s             // serialize model to plan file i.e. 'hrnet.engine'
  sudo ./hrnet -d  ../samples // deserialize plan file and run inference, the images in samples will be processed.
  

The test img:

Pytorch Result:

TRT Result: