Skip to content

Latest commit

 

History

History
66 lines (48 loc) · 2.12 KB

README.md

File metadata and controls

66 lines (48 loc) · 2.12 KB

LibTorch Inference

The LibTorch inference for yolort, both GPU and CPU are supported.

Dependencies

  • Ubuntu 18.04
  • LibTorch 1.8.0 / 1.9.0
  • TorchVision 0.9.0 / 0.10.0
  • OpenCV 3.4+
  • CUDA 10.2 [Optional]

We didn't impose too strong restrictions on the version of CUDA and Ubuntu systems.

Usage

  1. First, Setup the environment variables.

    export TORCH_PATH=$(dirname $(python -c "import torch; print(torch.__file__)"))
export LD_LIBRARY_PATH=$LD_LIBRARY_PATH:$TORCH_PATH/lib/

  2. Don't forget to compile TorchVision using the following scripts.

    git clone https://github.com/pytorch/vision.git
cd vision
git checkout release/0.9  # Double check the version of TorchVision currently in use
mkdir build && cd build
cmake .. -DTorch_DIR=$TORCH_PATH/share/cmake/Torch  # Set `-DWITH_CUDA=ON` if you're using GPU
make -j4
sudo make install

  3. Generate TorchScript model

    Unlike ultralytics's torch.jit.trace mechanism, We're using torch.jit.script to trace the YOLOv5 models which containing the whole pre-processing (especially with the letterbox ops) and post-processing (especially with the nms ops) procedures, as such you don't need to rewrite manually the C++ codes of pre-processing and post-processing.

    from yolort.models import yolov5n

model = yolov5n(pretrained=True)
model.eval()

traced_model = torch.jit.script(model)
traced_model.save("yolov5n.torchscript.pt")

  4. Then compile the source code.

    cd deployment/libtorch
mkdir build && cd build
cmake .. -DTorch_DIR=$TORCH_PATH/share/cmake/Torch
make

  5. Now, you can infer your own images.

    ./yolo_inference [--input_source ../../../test/assets/zidane.jpg]
                 [--checkpoint ../yolov5n.torchscript.pt]
                 [--labelmap ../../../notebooks/assets/coco.names]
                 [--gpu]  # GPU switch, which is optional, and set False as default