빠른 추론 라이브러리 tkDNN에서 자체 모델로 이동하는 방법

경위

yolov4의 추론을 고속화시킨 추론 라이브러리를 사용할 기회가 있었습니다만, 스스로 학습시킨 모델을 움직일 때에 조금 막혔으므로, 순서를 공유합니다. 이번에는 darknet의 학습이 끝난 모델(자전 모델)을 tkDNN로 사용하는 방법을 소개합니다.

tkDNN이란?

tkDNN은 yolo darknet을 가속화할 수 있는 추론 라이브러리입니다. 학습된 v3 또는 v4의 weights 파일과 config 파일, 클래스명의 파일을 준비하는 것만으로 사용할 수 있습니다. jetson이라면 대략 2배~3배 고속화할 수 있고, 양자화하면 4배 이상 빨라집니다. 덧붙여서 darknet 이외에도 ResNet이나 CenterNet, MobileNet등도 고속화할 수 있습니다.

절차

1. tkDNN 빌드

아래 단계에서 tkDNN을 빌드합니다. 환경 작성 등은 여기 를 참고해 주세요.

git clone https://github.com/ceccocats/tkDNN
cd tkDNN
mkdir build
cd build
cmake .. 
make
rm yolo4_fp32.rt
./test_yolo4

2. weights export

여기 의 리포지토리를 사용해, weights 파일을 export 합니다.

git clone https://git.hipert.unimore.it/fgatti/darknet.git
cd darknet
make
mkdir layers debug
./darknet export <path-to-cfg-file> <path-to-weights> layers

path-to-cfg-file에 darknet 학습 시 config 파일의 경로와 path-to-weiths에 학습된 모델의 경로를 넣습니다. 이것이 성공적으로 끝나면 debug와 layers에 bin 파일이 만들어져야합니다. debug와 layers 폴더를 tkDNN/build/yolo4에 넣으십시오. 이미 debug와 layers 폴더가 있다면 교체하십시오.

3. tkDNN 소스 코드 변경

기본 코드를 그대로 사용하면 이전 모델의 클래스나 설정을 반영할 수 없으므로 tkDNN 소스 코드를 편집합니다. 아래는 기본값으로 유지되지만 cfg_path와 name_path를 자신의 것으로 설정하십시오. names 파일은 학습한 클래스를 기입한 것을 스스로 작성해 주세요.

tkDNN/tests/darknet/yolo4.cpp

int main() {
    std::string bin_path  = "yolo4";
    std::vector<std::string> input_bins = { 
        bin_path + "/layers/input.bin"
    };
    std::vector<std::string> output_bins = {
        bin_path + "/debug/layer139_out.bin",
        bin_path + "/debug/layer150_out.bin",
        bin_path + "/debug/layer161_out.bin"
    };
    std::string wgs_path  = bin_path + "/layers";
    std::string cfg_path  = std::string(TKDNN_PATH) + "/tests/darknet/cfg/yolo4.cfg";
    std::string name_path = std::string(TKDNN_PATH) + "/tests/darknet/names/coco.names";
    downloadWeightsifDoNotExist(input_bins[0], bin_path, "https://cloud.hipert.unimore.it/s/d97CFzYqCPCp5Hg/download");

    ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
}

4. tkdnn 빌드 & weighs 파일 변환

아래 명령을 실행합니다. fp16으로 양자화하고 싶은 경우는 export TKDNN_MODE=FP16를 실행하면 괜찮습니다.

cd build
cmake ..
make
rm yolo4_fp32.rt
## (量子化)export TKDNN_MODE=FP16
./test_yolo4

완료하면, yolo4_fp32.rt(양자화한 경우는 yolo4_fp16.rt)라는 파일이 작성되고 있다고 생각합니다. 이것으로 완료됩니다. 마지막으로 아래 명령을 실행하면 데모 동영상에서 시도할 수 있습니다.

./demo yolo4_fp32.rt

요약

tkDNN을 만든 사람은 대단해! 리포지토리에서의 issue도 상당히 대답해 줍니다.

실수나 질문, 의견등 있으면 부담없이 코멘트해 주세요. 열심히 대답하니까(웃음).