[Pose Estimation] NVIDIA TLT의 BodyPoseNet

 

NVIDIA Transfer Learning Toolkit(이하 TLT)에서 "BodyPoseNet" 과 함께 COCO 데이터를 사용하여 2D Pose Estimation 모델을 학습하고 INT8 등으로 최적화하는 방법에 대해 다루었다. 이는 기존 오픈 소스인 OpenPose 보다 AP가 8% 정도 떨어지지만 속도면에서는 아주 월등히 우수함을 보여준다.

 

BodyPoseNet

 

NVIDIA TLT에서 제공하는 BodyPoseNet은 VGG와 같은 백본 네트워크를 이용하여 confidence map 및 PAF를 이용하여 픽셀 단위 예측을 수행한 후 multi stage refinement (0 ~ N stages)로 구성된 구조를 사용하는 fully connected convolution 모델이며, 18개의 키포인트들을 예측한다.

 

• nose, neck, right_shoulder, right_elbow, right_wrist, left_shoulder, left_elbow, left_wrist, right_hip, right_knee, right_ankle, left_hip, left_knee, left_ankle, right_eye, left_eye, right_ear, left_ear.

 

 

 

먼저 모델을 Pruning 하는 방법은 다음과 같다. 참고로 pth 변수는 pruning threshold 인 듯 한데, 0.05 에서 3.0 범위 사이의 값이 BodyPoseNet 모델을 pruning 하는데 정확도와 모델 크기의 trade off를 만족 시킬 수 있는 값이라고 한다. 

 

tlt bpnet prune -m $USER_EXPERIMENT_DIR/models/exp_m1_unpruned/bpnet_model.tlt \
                 -o $USER_EXPERIMENT_DIR/models/exp_m1_pruned/bpnet_model.pruned-0.05.tlt \
                 -eq union \
                 -pth 0.05 \
                 -k $KEY 

 

 

 

그 다음 pruning 을 수행한 모델은 이전의 가중치 중에서 유용한 정보를 제공하는 가중치가 소실되었을 수도 있기 때문에 정확도가 약간 떨어질 수 있다. 그러므로 정확도를 다시 확보하기 위해서 동일한 데이터 세트에 대해 pruning 된 모델을 다시 학습하는 것이 좋다. 

 

# Retraining using the pruned model as model graph
tlt bpnet train  -e $ SPECS_DIR / bpnet_retrain_m1_coco.yaml \
                 -r $ USER_EXPERIMENT_DIR / models / exp_m1_retrain \
                 -k $ KEY \
                 --gpus $ NUM_GPUS 

 

 

 

 

이렇게 다시 학습된 모델을 평가할 수 있는데, COCO 데이터 세트에 대해 평가한 결과는 다음과 같다.

 

Average Precision  (AP) @[ IoU=0.50:0.95 | area=   all | maxDets= 20 ] = 0.561
Average Precision  (AP) @[ IoU=0.50      | area=   all | maxDets= 20 ] = 0.776
Average Precision  (AP) @[ IoU=0.75      | area=   all | maxDets= 20 ] = 0.609
Average Precision  (AP) @[ IoU=0.50:0.95 | area=medium | maxDets= 20 ] = 0.567
Average Precision  (AP) @[ IoU=0.50:0.95 | area= large | maxDets= 20 ] = 0.556
... 

 

 

 

 

 

TLT에서는 이렇게 학습 및 최적화 된 모델을 export 하기 위해 다음과 같은 명령어를 지원한다. 

 

tlt bpnet export -m $ USER_EXPERIMENT_DIR / models / exp_m1_retrain / bpnet_model.tlt \
                  -e $ SPECS_DIR / bpnet_retrain_m1_coco.yaml \
                  -o $ USER_EXPERIMENT_DIR / models / exp_m1_final / bpnet_model.etlt \
                  -k $ KEY \
                  -t tfonnx 

 

 

 

 

 

 

BodyPoseNet 모델은 TensorRT 에서 INT8 inference mode를 지원하기 때문에 먼저 이를 위해 캘리브레이션을 수행해야한다. 캘리브레이션을 하기 위해서는 샘플 이미지들이 필요하다. TLT에서는 필요한 이미지들을 샘플링하는 스크립트를 아래와 같이 제공한다. 

 

# Number of calibration samples to use
 export NUM_CALIB_SAMPLES=2000
  
 python3 sample_calibration_images.py \
     -a $LOCAL_EXPERIMENT_DIR/data/annotations/person_keypoints_train2017.json \
     -i $LOCAL_EXPERIMENT_DIR/data/train2017/ \
     -o $LOCAL_EXPERIMENT_DIR/data/calibration_samples/ \
     -n $NUM_CALIB_SAMPLES \
     -pth 1 \
     --randomize 

 

 

 

 

 

 

그 다음 학습된 모델을 .etlt 형식으로 export 하고 INT8 캘리브레이션을 수행하여 TensorRT 엔진을 생성할 수 있다.

 

 

이 때 디렉토리(cal_image_dir)에 batch_size * batches 개의 이미지가 실제로 있는지 확인해야한다. 또한 INT8 모드가 아닌 FP16 모드로 엔진을 생성하려면 --data_type 을 FP16 으로 지정하면 된다. 

 

# Set dimensions of desired output model for inference/deployment
 export IN_HEIGHT=288
 export IN_WIDTH=384
 export IN_CHANNELS=3
 export INPUT_SHAPE=288x384x3
 # Set input name
 export INPUT_NAME=input_1:0
  
 tlt bpnet export \
     -m $USER_EXPERIMENT_DIR/models/exp_m1_retrain/bpnet_model.tlt \
     -o $USER_EXPERIMENT_DIR/models/exp_m1_final/bpnet_model.etlt \
     -k $KEY \
     -d $IN_HEIGHT,$IN_WIDTH,$IN_CHANNELS \
     -e $SPECS_DIR/bpnet_retrain_m1_coco.yaml \
     -t tfonnx \
     --data_type int8 \
     --engine_file $USER_EXPERIMENT_DIR/models/exp_m1_final/bpnet_model.$IN_HEIGHT.$IN_WIDTH.int8.engine \
     --cal_image_dir $USER_EXPERIMENT_DIR/data/calibration_samples/ \
     --cal_cache_file $USER_EXPERIMENT_DIR/models/exp_m1_final/calibration.$IN_HEIGHT.$IN_WIDTH.bin  \
     --cal_data_file $USER_EXPERIMENT_DIR/models/exp_m1_final/coco.$IN_HEIGHT.$IN_WIDTH.tensorfile \
     --batch_size 1 \
     --batches $NUM_CALIB_SAMPLES \
     --max_batch_size 1 \
     --data_format channels_last 

 

 

 

 

 

이렇게 생성된 엔진을 이용하여 어플리케이션에 적용해볼 수 있는데, 그 전에 고려해야할 것은 FP16 모드로 생성된 엔진과 FP32로 생성된 엔진이 거의 비슷한 정확도를 가지는지 확인해야하고, INT8 모델은 두 precision 모드와 2~3% AP 범위 내에 정확도를 가지는지 확인해야한다. 

 

그 다음 배포를 위한 모델을 export 하기 위한 명령어는 다음과 같다. 

 

tlt bpnet export
     -m $USER_EXPERIMENT_DIR/models/exp_m1_retrain/bpnet_model.tlt
     -o $USER_EXPERIMENT_DIR/models/exp_m1_final/bpnet_model.deploy.etlt
     -k $KEY
     -d $IN_HEIGHT,$IN_WIDTH,$IN_CHANNELS
     -e $SPECS_DIR/bpnet_retrain_m1_coco.txt
     -t tfonnx
     --data_type int8
     --cal_image_dir $USER_EXPERIMENT_DIR/data/calibration_samples/
     --cal_cache_file $USER_EXPERIMENT_DIR/models/exp_m1_final/calibration.$IN_HEIGHT.$IN_WIDTH.deploy.bin
     --cal_data_file $USER_EXPERIMENT_DIR/models/exp_m1_final/coco.$IN_HEIGHT.$IN_WIDTH.tensorfile
     --batch_size 1
     --batches $NUM_CALIB_SAMPLES
     --max_batch_size 1
     --data_format channels_last
     --engine_file $USER_EXPERIMENT_DIR/models/exp_m1_final/bpnet_model.$IN_HEIGHT.$IN_WIDTH.int8.deploy.engine
     --sdk_compatible_model 

 

 

 

 

 

모델의 네트워크 입력 해상도는 상향식 접근 방식의 정확도를 결정할 수 있는 주요 요인 중 하나인데, 해상도가 높을 수록 더 나은 정확도를 제공하게 된다. 그러나 입력 해상도가 높을 수록 CNN의 런타임도 높아지기 때문에 정확도와 런타임 간 trade off 를 고려하여, 해상도를 결정해야한다. 상대적으로 사이즈가 작은 사람들을 estimation 하기 위해서는 더 높은 네트워크 입력 높이를 사용하는 것이 좋다. 높이를 고정한 후 입력 데이터의 종횡비에 따라 너비를 결정할 수 있다.

 

이렇듯 다양한 해상도에 대한 정확도 및 런타임 변화에 대한 실험 결과는 다음과 같다. 

 

 

NVIDIA에서 제공하는 BodyPoseNet은 유명한 OpenPose와 비교하면 아래와 같은 성능을 가진다. AP는 약 8% 정도 낮은 결과를 보여준다. 

 

 

 

하지만 속도 측면에 있어서 어마어마한 차이를 보여준다. 

 

 

 

 

다음 표는 기본 매개변수들을 사용하여 TLT로 학습된 BodyPoseNet 모델의 추론 성능을 나타낸다. trtexec를 이용하여 모델 추론을 프로파일링 한 결과이다. 확실히 경험상 trtexec 모듈은 여러 GPU에서 TensorRT 엔진을 실험할 때 정말 편한 것 같다. 

 

 

 

 

NVIDIA TLT 에서는 Pose Estimation 모델 뿐만 아니라 다른 산업 분야에서도 잘 쓰일 수 있는 여러 모델들을 제공하고 있고, INT8 이라는 아주 매력적인 선택지를 제공하고 있기 때문에 앞으로 AI 업계에서 가장 발빠른 적응은 물론 다른 기업들과 달리 미래를 보고 앞서 나갈 수 있는 독보적인 위치에 있지 않을까 생각된다. 앞으로의 NVIDIA의 행보가 기대된다. 

 

 

 

 

 

 

 

참고자료 1 : https://developer.nvidia.com/blog/training-optimizing-2d-pose-estimation-model-with-tlt-part-2/

Training and Optimizing a 2D Pose Estimation Model with the NVIDIA Transfer Learning Toolkit, Part 2 | NVIDIA Developer Blog

 

참고자료 2 : https://docs.nvidia.com/tlt/tlt-user-guide/text/bodypose_estimation/bodyposenet.html

Body Pose Estimation — Transfer Learning Toolkit 3.0 documentation

 

참고자료 3 : https://ngc.nvidia.com/catalog/models/nvidia:tlt_bodyposenet

NVIDIA NGC