如何使用COCO數(shù)據(jù)集訓(xùn)練和優(yōu)化二維姿態(tài)估計模型

存在開發(fā)姿態(tài)估計的開源方法，但在推理性能方面不是最優(yōu)的，并且集成到生產(chǎn)應(yīng)用程序中非常耗時。通過這篇文章，我們將向您展示如何開發(fā)和部署姿勢估計模型，這些模型易于跨設(shè)備配置文件使用，性能非常好，并且非常精確。

姿勢估計已與 NVIDIA 轉(zhuǎn)移學(xué)習(xí)工具包（ TLT ） 3.0 集成，因此您可以利用 TLT 的所有功能（如模型修剪和量化）來創(chuàng)建精確和高性能的模型。經(jīng)過訓(xùn)練后，您可以部署此模型進(jìn)行推理以獲得實時性能。

本系列文章將引導(dǎo)您完成培訓(xùn)、優(yōu)化、部署實時高性能姿勢估計模型的步驟。在第 1 部分中，您將學(xué)習(xí)如何使用開放源代碼 COCO 數(shù)據(jù)集訓(xùn)練二維姿勢估計模型。在 第 2 部分 中，您將學(xué)習(xí)如何為推理吞吐量優(yōu)化模型，然后使用 TLT-CV 推理管道部署模型。我們將 TLT 訓(xùn)練模型與其他最先進(jìn)的模型進(jìn)行了比較。

用 TLT 訓(xùn)練二維姿態(tài)估計模型

在本節(jié)中，我們將介紹有關(guān)使用 TLT 訓(xùn)練 2D 姿勢估計模型的以下主題：

Methodology

環(huán)境設(shè)置

數(shù)據(jù)準(zhǔn)備

實驗配置文件

Training

Evaluation

模型驗證

Methodology

BodyPoseNet 模型的目標(biāo)是預(yù)測給定輸入圖像中每個人的骨架，骨架由關(guān)鍵點(diǎn)和關(guān)鍵點(diǎn)之間的連接組成。

兩種常用的姿態(tài)估計方法是自頂向下和自下而上。自頂向下的方法通常使用對象檢測網(wǎng)絡(luò)來定位幀中所有人的邊界框，然后使用姿勢網(wǎng)絡(luò)來定位該邊界框內(nèi)的身體部位。顧名思義，自下而上的方法從下到上構(gòu)建骨架。它首先檢測一個框架內(nèi)的所有人體部位，然后使用一種方法對屬于特定人的部位進(jìn)行分組。

采用自下而上的方法有幾個原因。一是推理性能較高。與自頂向下的姿勢估計方法不同，自下而上的方法不需要單獨(dú)的人檢測器。計算不會隨場景中的人數(shù)線性縮放。這使您能夠?qū)崿F(xiàn)擁擠場景的實時性能。此外，自底向上還具有全局上下文的優(yōu)點(diǎn)，因為將整個圖像作為輸入提供給網(wǎng)絡(luò)。它可以更好地處理復(fù)雜的姿勢和擁擠。

鑒于這些原因，這種方法的目的是實現(xiàn)有效的單桿，自下而上的姿態(tài)估計，同時也提供競爭力的準(zhǔn)確性。本文使用的默認(rèn)模型是完全卷積模型，由主干網(wǎng)、初始預(yù)測階段組成，該階段對置信圖（ heatmap ）和部分親和場（ PAF ）進(jìn)行像素級預(yù)測，然后對初始預(yù)測進(jìn)行多級細(xì)化（ 0 到 N 階段）。此解決方案簡化并抽象了自底向上方法的許多復(fù)雜性，同時允許針對特定應(yīng)用程序調(diào)整必要的旋鈕。

圖 1 。默認(rèn)模型架構(gòu)的簡化框圖。

PAFs 是一種用自下而上的方法表示關(guān)聯(lián)分?jǐn)?shù)的方法。有關(guān)詳細(xì)信息，請參閱 基于部分相似域的實時多人二維姿態(tài)估計 。它由一組二維向量場組成，對肢體的位置和方向進(jìn)行編碼。這與熱圖相關(guān)聯(lián)，用于在后處理期間通過執(zhí)行二部匹配和關(guān)聯(lián)身體部位候選來構(gòu)建骨架。

環(huán)境設(shè)置

NVIDIA TLT 工具包有助于抽象出 AI / DL 框架的復(fù)雜性，并使您能夠更快地構(gòu)建生產(chǎn)質(zhì)量模型，而無需編碼。有關(guān)硬件和軟件要求、設(shè)置所需依賴項以及安裝 TLT 啟動器的更多信息，請參閱 TLT 快速入門指南 。

使用以下命令下載最新示例：

ngc registry resource download-version "nvidia/tlt_cv_samples:v1.1.0"

您可以在tlt_cv_samples:v1.1.0/bpnet找到示例筆記本，其中還包含所有詳細(xì)步驟。

為 cleaner 命令行命令設(shè)置 env 變量。更新以下變量值：

 export KEY=
 export NUM_GPUS=1
 # Local paths
 # The dataset is expected to be present in $LOCAL_PROJECT_DIR/bpnet/data. 
 export LOCAL_PROJECT_DIR=/home//tlt-experiments
 export SAMPLES_DIR=/home//tlt_cv_samples_vv1.1.0
 # Container paths
 export USER_EXPERIMENT_DIR=/workspace/tlt-experiments/bpnet
 export DATA_DIR=/workspace/tlt-experiments/bpnet/data
 export SPECS_DIR=/workspace/examples/bpnet/specs
 export DATA_POSE_SPECS_DIR=/workspace/examples/bpnet/data_pose_config
 export MODEL_POSE_SPECS_DIR=/workspace/examples/bpnet/model_pose_config 

要運(yùn)行 TLT 啟動程序，請使用~/.tlt_mounts.json文件將本地計算機(jī)上的~/ TLT-experiments 目錄映射到 Docker 容器。有關(guān)詳細(xì)信息，請參閱TLT 發(fā)射器。

創(chuàng)建~/.tlt_mounts.json文件并更新其中的以下內(nèi)容：

 {
  "Mounts": [
  {
  "source": "/home//tlt-experiments",
  "destination": "/workspace/tlt-experiments"
  },
  {
  "source": "/home//tlt_cv_samples_vv1.1.0/bpnet/specs",
  "destination": "/workspace/examples/bpnet/specs"
  },
  {
  "source": "/home//tlt_cv_samples_vv1.1.0/bpnet/data_pose_config",
  "destination": "/workspace/examples/bpnet/data_pose_config"
  },
  {
  "source": "/home//tlt_cv_samples_vv1.1.0/bpnet/model_pose_config",
  "destination": "/workspace/examples/bpnet/model_pose_config"
  }
  ]
 } 

確保要裝載的源目錄路徑有效。這會將主機(jī)上的路徑/home//tlt-experiments裝載為容器內(nèi)的路徑/workspace/tlt-experiments。它還將下載的規(guī)范裝載到主機(jī)上，使其成為容器內(nèi)的路徑/workspace/examples/bpnet/specs、/workspace/examples/bpnet/data_pose_config和/workspace/examples/bpnet/model_pose_config。

通過運(yùn)行以下命令，確保已安裝所需的依賴項：

 # Install requirements 
 pip3 install -r $SAMPLES_DIR/deps/requirements-pip.txt 

下載預(yù)訓(xùn)練模型

首先，設(shè)置一個NGC帳戶，然后下載預(yù)訓(xùn)練模型。目前，只支持 vgg19 主干網(wǎng)。

 # Create the target destination to download the model.
 mkdir -p $LOCAL_EXPERIMENT_DIR/pretrained_vgg19/
 
 # Download the pretrained model from NGC
 ngc registry model download-version nvidia/tlt_bodyposenet:vgg19 \
  --dest $LOCAL_EXPERIMENT_DIR/pretrained_vgg19 

數(shù)據(jù)準(zhǔn)備

我們以本文中的 COCO （ context 上的公共對象） 2017 數(shù)據(jù)集為例。下載數(shù)據(jù)集并按照說明提?。?/p>

• train_data
• val_data
• annotations

將 images 目錄解壓到$LOCAL_DATA_DIR目錄中，并將注釋解壓到$LOCAL_DATA_DIR/annotations中。

要準(zhǔn)備用于訓(xùn)練的數(shù)據(jù)，必須生成分段掩碼，用于掩蓋未標(biāo)記人員和 TFR 記錄的丟失，以提供給訓(xùn)練管道。掩碼文件夾基于coco_spec.json文件中提供的路徑[mask_root_dir_path目錄是root_directory_path的相對路徑，mask_root_dir_path和annotation_root_dir_path也是如此。

 # Generate TFRecords for training dataset
 tlt bpnet dataset_convert \
  -m 'train' \
  -o $DATA_DIR/train \
  --generate_masks \
  --dataset_spec $DATA_POSE_SPECS_DIR/coco_spec.json
 
 # Generate TFRecords for validation dataset
 tlt bpnet dataset_convert \
  -m 'test' \
  -o $DATA_DIR/val \
  --generate_masks \
  --dataset_spec $DATA_POSE_SPECS_DIR/coco_spec.json 

要將此示例用于自定義數(shù)據(jù)集，請執(zhí)行以下操作：

• 以類似于 COCO 數(shù)據(jù)集的格式準(zhǔn)備數(shù)據(jù)和注釋。
• 在 data \ u pose \ u config 下創(chuàng)建一個數(shù)據(jù)集規(guī)范，類似于 coco \ u spec . json ，其中包括數(shù)據(jù)集路徑、姿勢配置、遮擋標(biāo)記約定等。
• 將注釋轉(zhuǎn)換為 COCO 注釋格式。

有關(guān)更多信息，請參閱以下文檔：

• COCO 數(shù)據(jù)格式
• 數(shù)據(jù)格式說明
• 為數(shù)據(jù)集轉(zhuǎn)換器創(chuàng)建配置文件

列車試驗配置文件

下一步是為培訓(xùn)配置 spec 文件。實驗規(guī)范文件是必不可少的，因為它編譯了實現(xiàn)良好模型所需的所有超參數(shù)。 BodyPoseNet 訓(xùn)練的規(guī)范文件配置訓(xùn)練管道的以下組件：

• Trainer
• Dataloader
• Augmentation
• 標(biāo)簽處理機(jī)
• Model
• Optimizer

您可以在$SPECS_DIR/bpnet_train_m1_coco.yaml中找到默認(rèn)規(guī)范文件。我們在規(guī)范文件的每個組件上展開，但這里不包括所有參數(shù)。有關(guān)詳細(xì)信息，請參閱創(chuàng)建列車試驗配置文件。

培訓(xùn)師（頂級配置）

頂層實驗配置包括實驗的基本參數(shù)；例如，歷元數(shù)、預(yù)訓(xùn)練權(quán)重、是否加載預(yù)訓(xùn)練圖等。根據(jù)checkpoint_n_epoch值保存加密的檢查點(diǎn)。下面是一些頂級配置的代碼示例。

checkpoint_dir: /workspace/tlt-experiments/bpnet/models/exp_m1_unpruned
 checkpoint_n_epoch: 5
 num_epoch: 100
 pretrained_weights: /workspace/tlt-experiments/bpnet/pretrained_vgg19/tlt_bodyposenet_vvgg19/vgg_19.hdf5
 load_graph: False
 use_stagewise_lr_multipliers: True
 ... 

所有路徑（checkpoint_dir和pretrained_weights）都是 Docker 容器的內(nèi)部路徑。要驗證正確性，請選中~/.tlt_mounts.json。有關(guān)這些參數(shù)的更多信息，請參閱身體姿勢訓(xùn)練器部分。

Dataloader

本節(jié)幫助您定義數(shù)據(jù)路徑、圖像配置、目標(biāo)姿勢配置、規(guī)范化參數(shù)等。augmentation_config部分提供了一些動態(tài)增強(qiáng)選項。它支持基本的空間增強(qiáng)，例如翻轉(zhuǎn)、縮放、旋轉(zhuǎn)和平移，這些都可以在訓(xùn)練實驗之前配置。label_processor_config部分提供了配置地面實況要素圖生成所需的參數(shù)。

dataloader:
  batch_size: 10
  pose_config:
  target_shape: [32, 32]
  pose_config_path: /workspace/examples/bpnet/model_pose_config/bpnet_18joints.json
  image_config:
  image_dims:
  height: 256
  width: 256
  channels: 3
  image_encoding: jpg
  dataset_config:
  root_data_path: /workspace/tlt-experiments/bpnet/data/
  train_records_folder_path: /workspace/tlt-experiments/bpnet/data
  train_records_path: [train-fold-000-of-001]
  dataset_specs:
  coco: /workspace/examples/bpnet/data_pose_config/coco_spec.json
  normalization_params:
  ...
  augmentation_config:
  spatial_augmentation_mode: person_centric
  spatial_aug_params:
  flip_lr_prob: 0.5
  flip_tb_prob: 0.0
  ...
  label_processor_config:
  paf_gaussian_sigma: 0.03
  heatmap_gaussian_sigma: 7.0
  paf_ortho_dist_thresh: 1.0 

• target_shape值取決于image_dims和模型步幅值（target_shape=input_shape/model stride）。當(dāng)前模型的步幅為 8 。
• 確保使用與dataset_spec中的root_directory_path相同的root_data_path值。dataset_spec中的掩碼和圖像數(shù)據(jù)目錄相對于root_data_path。
• 所有路徑，包括pose_config_path、dataset_config和dataset_specs，都是 Docker 的內(nèi)部路徑。
• 支持多種spatial_augmentation_modes：
  • person_centric：增強(qiáng)是圍繞一個對基本真相感興趣的人。
  • standard：增強(qiáng)是標(biāo)準(zhǔn)的（即，以圖像中心為中心），并且保留圖像的縱橫比。
  • standard_with_fixed_aspect_ratio：與標(biāo)準(zhǔn)相同，但縱橫比固定為網(wǎng)絡(luò)輸入縱橫比。

有關(guān)每個參數(shù)的詳細(xì)信息，請參閱Dataloader部分。

Model

BodyPoseNet 模型可以使用 spec 文件中的 model 選項進(jìn)行配置。下面是一個示例模型配置，用于實例化基于 VGG19 主干網(wǎng)的自定義模型。

model:
 backbone_attributes:
  architecture: vgg
 stages: 3
 heat_channels: 19
 paf_channels: 38
 use_self_attention: False
 data_format: channels_last
 use_bias: True
 regularization_type: l1
 kernel_regularization_factor: 5.0e-4
 bias_regularization_factor: 0.0
 ... 

網(wǎng)絡(luò)中用于姿勢估計的總階段數(shù)（細(xì)化階段+ 1 ）由stages參數(shù)捕獲，該參數(shù)取任何值>= 2 。我們建議在修剪前訓(xùn)練網(wǎng)絡(luò)時使用 L1 正則化器，因為 L1 正則化使修剪網(wǎng)絡(luò)權(quán)重更容易。有關(guān)模型中每個參數(shù)的詳細(xì)信息，請參閱Model部分。

優(yōu)化

本節(jié)介紹如何配置優(yōu)化器和學(xué)習(xí)速率計劃：

optimizer:
  __class_name__: WeightedMomentumOptimizer
  learning_rate_schedule:
  __class_name__: SoftstartAnnealingLearningRateSchedule
  soft_start: 0.05
  annealing: 0.5
  base_learning_rate: 2.e-5
  min_learning_rate: 8.e-08
  momentum: 0.9
  use_nesterov: False 

默認(rèn)的base_learning_rate是為單個 GPU 訓(xùn)練設(shè)置的。要使用多 GPU 訓(xùn)練，可能需要修改“學(xué)習(xí)率”值以獲得類似的精度。在大多數(shù)情況下，將學(xué)習(xí)率提高一倍$NUM_GPUS將是一個良好的開端。例如，如果您使用兩個 GPU ，請在一個 GPU 設(shè)置中使用2 * base_learning_rate，如果您使用四個 GPU ，請使用4 * base_learning_rate。有關(guān)模型中每個參數(shù)的詳細(xì)信息，請參閱Optimizer部分。

訓(xùn)練

在完成生成 TFRecords 和 mask 的步驟并設(shè)置了一個 train 規(guī)范文件之后，現(xiàn)在就可以開始訓(xùn)練 body pose estimation 網(wǎng)絡(luò)了。使用以下命令啟動培訓(xùn)：

tlt bpnet train -e $SPECS_DIR/bpnet_train_m1_coco.yaml \
  -r $USER_EXPERIMENT_DIR/models/exp_m1_unpruned \
  -k $KEY \
  --gpus $NUM_GPUS 

使用更多 GPU 進(jìn)行培訓(xùn)可以使網(wǎng)絡(luò)更快地接收更多數(shù)據(jù)，從而在開發(fā)過程中節(jié)省寶貴的時間。 TLT 支持多 GPU 訓(xùn)練，因此可以使用多個 GPU 并行訓(xùn)練模型。我們建議使用四個或更多的 GPU 來訓(xùn)練模型，因為一個 GPU MIG ht 需要幾天才能完成。訓(xùn)練時間大致減少了一個系數(shù)$NUM_GPUS。確保根據(jù)Optimizer一節(jié)中描述的線性縮放方法相應(yīng)地更新學(xué)習(xí)速率。

BodyPoseNet 支持從檢查點(diǎn)重新啟動。如果訓(xùn)練作業(yè)過早終止，只需重新運(yùn)行相同的命令，就可以從上次保存的檢查點(diǎn)恢復(fù)訓(xùn)練。重新啟動培訓(xùn)時，請確保使用相同數(shù)量的 GPU 。

評估

從配置推斷和評估規(guī)范文件開始。下面的代碼示例是一個示例規(guī)范：

model_path: /workspace/tlt-experiments/bpnet/models/exp_m1_unpruned/bpnet_model.tlt
 train_spec: /workspace/examples/bpnet/specs/bpnet_train_m1_coco.yaml
 input_shape: [368, 368]
 # choose from: {pad_image_input, adjust_network_input, None}
 keep_aspect_ratio_mode: adjust_network_input
 output_stage_to_use: null
 output_upsampling_factor: [8, 8]
 heatmap_threshold: 0.1
 paf_threshold: 0.05
 multi_scale_inference: False
 scales: [0.5, 1.0, 1.5, 2.0] 

此處的input_shape值可以不同于用于培訓(xùn)的input_dims值。multi_scale_inference參數(shù)可在提供的比例上啟用多比例優(yōu)化。因為您使用的是步幅 8 的模型，output_upsampling_factor設(shè)置為 8 。

為了使評估與自底向上的人體姿勢估計研究保持一致，有兩種模式和規(guī)范文件來評估模型：

• $SPECS_DIR/infer_spec.yaml：單刻度，非嚴(yán)格輸入。此配置對輸入圖像進(jìn)行單比例推斷。通過固定網(wǎng)絡(luò)輸入的一側(cè)（高度或?qū)挾龋?，并調(diào)整另一側(cè)以匹配輸入圖像的縱橫比，來保持輸入圖像的縱橫比。
• $SPECS_DIR/infer_spec_refine.yaml：多尺度、非嚴(yán)格輸入。此配置對輸入圖像進(jìn)行多尺度推斷。電子秤是可配置的。

還有一種模式主要用于驗證最終導(dǎo)出的 TRT 模型。在后面的章節(jié)中使用這個。

• $SPECS_DIR/infer_spec_strict.yaml：單刻度，嚴(yán)格輸入。此配置對輸入圖像進(jìn)行單比例推斷。當(dāng) TRT 模型的輸入尺寸固定時，根據(jù)需要在側(cè)面填充圖像以適應(yīng)網(wǎng)絡(luò)輸入尺寸，從而保留輸入圖像的縱橫比。

--model_filename參數(shù)重寫推理規(guī)范文件中的model_path變量。

要計算模型，請使用以下命令：

# Single-scale evaluation
 tlt bpnet evaluate --inference_spec $SPECS_DIR/infer_spec.yaml \
  --model_filename $USER_EXPERIMENT_DIR/models/exp_m1_unpruned/$MODEL_CHECKPOINT \
  --dataset_spec $DATA_POSE_SPECS_DIR/coco_spec.json \
  --results_dir $USER_EXPERIMENT_DIR/results/exp_m1_unpruned/eval_default \
  -k $KEY 

模型驗證

現(xiàn)在您已經(jīng)訓(xùn)練了模型，運(yùn)行推斷并驗證預(yù)測。要使用 TLT 直觀地驗證模型，請使用tlt bpnet inference命令。該工具支持對.tlt模型和 TensorRT.engine模型運(yùn)行推理。它在detections.json中生成帶注釋的圖像，在這些圖像上渲染骨架，并逐幀序列化關(guān)鍵點(diǎn)標(biāo)簽和元數(shù)據(jù)。例如，要使用經(jīng)過訓(xùn)練的.tlt模型運(yùn)行推理，請運(yùn)行以下命令：

tlt bpnet inference --inference_spec $SPECS_DIR/infer_spec.yaml \
  --model_filename $USER_EXPERIMENT_DIR/models/exp_m1_unpruned/$MODEL_CHECKPOINT \
  --input_type dir \
  --input $USER_EXPERIMENT_DIR/data/sample_images \
  --results_dir $USER_EXPERIMENT_DIR/results/exp_m1_unpruned/infer_default \
  --dump_visualizations \
  -k $KEY 

圖 1 顯示了原始圖像的一個示例，圖 2 顯示了渲染姿勢結(jié)果的輸出圖像。如您所見，該模型對不同于 COCO 訓(xùn)練數(shù)據(jù)的圖像具有魯棒性。

結(jié)論

在這篇文章中，您學(xué)習(xí)了如何使用 TLT 中的 BodyPoseNet 應(yīng)用程序訓(xùn)練身體姿勢模型。這篇文章展示了從 NGC 獲取一個帶有預(yù)訓(xùn)練主干的開源 COCO 數(shù)據(jù)集，用 TLT 訓(xùn)練一個模型。要優(yōu)化用于推理和部署的訓(xùn)練模型，請參見 二維姿態(tài)估計模型的訓(xùn)練與優(yōu)化，第 2 部分 。

關(guān)于作者

Sakthivel Sivaraman 是 NVIDIA 的高級軟件工程師，專注于開發(fā)深度學(xué)習(xí)和計算機(jī)視覺解決方案，并將其部署到邊緣。在 2018 參加“ ZVK3］之前，他從賓夕法尼亞大學(xué)獲得機(jī)器人學(xué)博士學(xué)位。他的研究興趣包括計算機(jī)視覺、深度學(xué)習(xí)和機(jī)器人技術(shù)。

Rajath Shetty 是 NVIDIA 的工程經(jīng)理，負(fù)責(zé)在汽車和醫(yī)療保健領(lǐng)域應(yīng)用深度學(xué)習(xí)和計算機(jī)視覺的項目。他的興趣涉及邊緣計算、算法和人工智能應(yīng)用軟件棧。他擁有喬治亞理工學(xué)院的電子和計算機(jī)工程碩士學(xué)位。

Chintan Shah 是 NVIDIA 的產(chǎn)品經(jīng)理，專注于智能視頻分析解決方案的 AI 產(chǎn)品。他管理工具箱，用于有效的深度學(xué)習(xí)培訓(xùn)和實時推理。在他之前的工作中，他正在為 NVIDIA GPU 開發(fā)硬件 IP 。他擁有北卡羅來納州立大學(xué)電氣工程碩士學(xué)位。

Niral Pathak 是 NVIDIA 的軟件工程師，致力于計算機(jī)視覺解決方案的部署。他擁有加州大學(xué)圣地亞哥分校電子和計算機(jī)工程碩士學(xué)位。他的興趣包括計算機(jī)視覺、深度學(xué)習(xí)和機(jī)器人技術(shù)。

審核編輯：郭婷