圖像分類 (image classification) 是視覺人工智能的最基礎(chǔ)應(yīng)用，目前 TAO 模型訓(xùn)練工具支持 resnet、vgg、mobilenet_v1、mobilenet_v2、googlenet、cspdarknet、darknet、efficientnet_b0、efficientnet_b1、cspdarknet_tiny 等 10 種神經(jīng)網(wǎng)絡(luò)，其中 resnet、vgg、darknet、cspdarknet 還有不同結(jié)構(gòu)層數(shù)的區(qū)分。

在 TAO 啟動(dòng)器中只用 “classification” 這個(gè)任務(wù)指令，去面對(duì)上面所列出的十多種網(wǎng)絡(luò)結(jié)構(gòu)，其關(guān)鍵處就是結(jié)合“配置文件”的協(xié)助，因此這個(gè)配置文件的設(shè)置組是 TAO 工具的核心所在，也是本文一開始要花時(shí)間說(shuō)明的部分。

在 cv_samples 下面的 classification 項(xiàng)目，是 TAO 工具的圖像分類模型訓(xùn)練的標(biāo)準(zhǔn)范例，里面提供 classification.ipynb 執(zhí)行腳本與 specs 目錄下兩個(gè) .cfg 配置文件，基本上在執(zhí)行腳本里只需要修改一些路徑的設(shè)定就行，其余細(xì)節(jié)都在配置文件里面進(jìn)行調(diào)整，包括所選擇的神經(jīng)網(wǎng)絡(luò)種類與結(jié)構(gòu)。

這個(gè)項(xiàng)目以 PascalVOC 2012 數(shù)據(jù)集作為訓(xùn)練源，與 TAO 所支持的格式并不相同，因此在訓(xùn)練之前還需要進(jìn)行格式的轉(zhuǎn)換。還好這里的轉(zhuǎn)換比較簡(jiǎn)單，只要將圖像分別存放在以“類別”所命名的目錄下就可以，同時(shí)分割成訓(xùn)練與校驗(yàn)兩個(gè)不同用途的數(shù)據(jù)集，在腳本里透過(guò)一段簡(jiǎn)單的 Python 代碼就能完成，并沒有什么難度。

最后在訓(xùn)練之前，可以選擇是否啟動(dòng) “遷移學(xué)習(xí)” 的功能?在這個(gè)實(shí)驗(yàn)中也會(huì)示范這中間所得到的精度差異，讓大家直接體驗(yàn)到遷移學(xué)習(xí)所帶來(lái)的的好處。

以下就將幾個(gè)執(zhí)行重點(diǎn)提出來(lái)說(shuō)明，協(xié)助大家可以輕松地執(zhí)行。

配置文件內(nèi)容

TAO 的 CLI 指令集只使用一個(gè) classification 去面對(duì)所有的圖形分類應(yīng)用，其余的工作就全部交給配置文件去處理。在圖像分類的配置文件里，主要有以下三個(gè)配置組：

1、model_config：存放神經(jīng)網(wǎng)絡(luò)種類、結(jié)構(gòu)(層數(shù))與特性的內(nèi)容，以下列出比較重要的部分：

(1) arch：網(wǎng)絡(luò)名稱，可使用所有已支持的網(wǎng)絡(luò)，包括 resnet、vgg、mobilenet_v1、mobilenet_v2、googlenet、cspdarknet、darknet、efficientnet_b0、efficientnet_b1、cspdarknet_tiny 等;

(2) n_layer：有些網(wǎng)絡(luò)有多種結(jié)構(gòu)層，例如 resnet有10/18/34/50/101、vgg 有16/19、darknet 有 19/53、cspdarknet 有 19/53 等;

其余參數(shù)都按照配置文件里面所設(shè)定的值，因?yàn)?NGC 提供的預(yù)訓(xùn)練模型是按照這些參數(shù)所訓(xùn)練的，能在遷移學(xué)習(xí)過(guò)程中得到比較好的效果。

2、train_config：執(zhí)行訓(xùn)練時(shí)所需要參考的

(1) train_dataset_path: 訓(xùn)練用數(shù)據(jù)集的位置，需要輸入在容器內(nèi)的完整路徑

(2) val_dataset_path: 校驗(yàn)用數(shù)據(jù)集的位置，需要輸入在容器內(nèi)的完整路徑

(3) pretrained_model_path: 預(yù)訓(xùn)練模型位置，需要輸入在容器內(nèi)的完整路徑

(4) batch_size_per_gpu：請(qǐng)根據(jù)GPU顯存大小進(jìn)行調(diào)整

(5) n_epochs：訓(xùn)練的回合數(shù)

(6) n_workers：CPU的并行線程數(shù)，請(qǐng)根據(jù)實(shí)際CPU核數(shù)量進(jìn)行調(diào)整

其余參數(shù)清先安裝范例所提供的設(shè)置。

3、eval_config：評(píng)估模型時(shí)所需要用到的參數(shù)，需要輸入在容器內(nèi)的完整路徑

(1) eval_dataset_path: 測(cè)試用數(shù)據(jù)集路徑，與前面相同

(2) model_path: 在前面訓(xùn)練模型過(guò)程中所生成的 .tlt 模型文件中，挑選效果最好的一個(gè)模型來(lái)進(jìn)行評(píng)估，通常最后一個(gè)的效果會(huì)最好。這里的 “resnet_080.tlt” 是因?yàn)橛?xùn)練回合數(shù)為 80。

其余參數(shù)請(qǐng)先按照范例所提供的設(shè)置。

以上關(guān)于路徑部分的設(shè)置，只要遵循腳本一開始 “0. Set up env variables and map drives” 的配置規(guī)則，文件里基本上不需要做修改。其他參數(shù)的定義與設(shè)定值，請(qǐng)?jiān)L問 https://docs.nvidia.com/tao/tao-toolkit/text/image_classification.html#model-config 有完整的說(shuō)明。

數(shù)據(jù)集格式轉(zhuǎn)換與分配

TAO 支持的圖像分類數(shù)據(jù)格式，是以“分類名”作為路徑名，將所有該類的圖像都放置到分類名目錄之下，例如有 20 個(gè)分類的數(shù)據(jù)集，就會(huì)有 20 個(gè)分類目錄。

本實(shí)驗(yàn)使用 PascalVOC 的 VOCtrainval_11-May-2012.tar (1.9GB) 數(shù)據(jù)集，可以用腳本提供的鏈接，也可以在https://pan.baidu.com/s/1JhnBCRi32xblhSSWFmF5aA (密碼: gg95) 下載壓縮文件。

請(qǐng)按照腳本的數(shù)據(jù)集處理過(guò)程，包括放置路徑與解壓縮的指令。由于這個(gè)數(shù)據(jù)集的格式并不符合 TAO 的要求，因此腳本中使用一段 Python 代碼來(lái)進(jìn)行轉(zhuǎn)換，詳細(xì)內(nèi)容請(qǐng)自行閱讀代碼，其主要功能就是將下圖左邊 VOC 路徑結(jié)構(gòu)轉(zhuǎn)換成右邊符合 TAO 要求的結(jié)構(gòu)，存放在 “formated” 目錄下。

接著再將 “formated” 的數(shù)據(jù)，隨機(jī)分割成 train、val與test 三大類存放到 “split” 目錄下，作為后面執(zhí)行訓(xùn)練、校驗(yàn)與測(cè)試時(shí)所需要的數(shù)據(jù)路徑。

動(dòng)遷移學(xué)習(xí)功能

TAO 的模型訓(xùn)練繼承 TLT 的遷移學(xué)習(xí)功能，只要在配置文件中的“train_config” 配置組里，指定好 “pretrained_model_path:” 的路徑就可以，如果不想啟用這個(gè)功能，只要在這個(gè)參數(shù)前面用 “#” 去關(guān)閉就行。

在 NGC 提供 300 多個(gè)預(yù)訓(xùn)練模型，我們需要挑選合適的模型來(lái)協(xié)助執(zhí)行遷移學(xué)習(xí)，這部分的下載需要使用 NGC 的指令來(lái)操作，在腳本中已經(jīng)提供完整的安裝步驟，只要執(zhí)行命令塊就行，然后用以下指令去下載所選的模型：

ngc registry model download-version 模型全名 --dest 本地存放路徑

這里配合所使用的網(wǎng)絡(luò)種類，挑選 nvidia/tao/pretrained_classification:resnet18 作為本次的預(yù)訓(xùn)練模型，文件大小在 80MB 左右。

下圖左方是這個(gè)實(shí)驗(yàn)中，關(guān)閉遷移學(xué)習(xí)功能所訓(xùn)練出來(lái)的模型，經(jīng)過(guò)下面評(píng)估步驟所測(cè)算出來(lái)的結(jié)果，下圖右方則是在基于 NGC 下載的預(yù)訓(xùn)練模型基礎(chǔ)上，利用遷移學(xué)習(xí)功能所訓(xùn)練出來(lái)的模型效果，總體來(lái)說(shuō)得到的精準(zhǔn)度提升還是很明顯的。

但如果是缺乏合適預(yù)訓(xùn)練模型的狀況下，也不能硬是挑選不合適的模型來(lái)做遷移學(xué)習(xí)，可能得到的效果會(huì)更差。

執(zhí)行模型訓(xùn)練

當(dāng)前面的準(zhǔn)備工作都做到位之后，這個(gè)步驟就是水到渠成。執(zhí)行訓(xùn)練所耗費(fèi)的時(shí)間就只跟 GPU 卡的總體計(jì)算資源有關(guān)系，如果您的設(shè)備有一張以上 GPU 計(jì)算卡的話，可以透過(guò) TAO 指令以下參數(shù)，去調(diào)整可用的計(jì)算資源：

--gpus [N]：指定要調(diào)用的 GPU 卡數(shù)量;

--gpu_index 編號(hào)1 [編號(hào)2...]

例如設(shè)備上有 2 片 GPU 卡時(shí)，想要全部用上的話就直接用 “--gpus 2” 參數(shù)進(jìn)行調(diào)用;如果設(shè)備上有 4 片 GPU 計(jì)算卡，想要指定調(diào)用 0 號(hào)與 3 號(hào)這兩張卡的時(shí)，就用 “--gpus 2--gpu_index 0 3” 參數(shù)，就能精準(zhǔn)調(diào)用指定的GPU來(lái)進(jìn)行訓(xùn)練。

在這個(gè)圖像分類實(shí)驗(yàn)中，我們以 NVIDIA RTX2070 與 RTX3070 進(jìn)行測(cè)試，這兩張計(jì)算卡都具有 8GB 顯存。在固定回合數(shù)為 80、batch_size_per_gpu 為 64 的條件下，所耗費(fèi)的訓(xùn)練時(shí)間如下：

單GPU/3070-8G : 52分19秒

單GPU/2070-8G : 62分44秒

雙GPU/3070+2070 : 38分03秒

修剪模型

這個(gè)步驟的目的，是要在維持足夠精準(zhǔn)度的前提下縮小模型的尺寸，越小的模型會(huì)消耗更少的計(jì)算資源，特別是內(nèi)存與顯存。

這對(duì)于計(jì)算資源充沛的計(jì)算設(shè)備來(lái)說(shuō)，所體現(xiàn)的差異并不明顯，但是對(duì)于計(jì)算資源相對(duì)吃緊的邊緣設(shè)備來(lái)說(shuō)，就有非常大的影響，甚至影響推理識(shí)別的性能，因此請(qǐng)自行決定是否需要執(zhí)行這個(gè)步驟。

這里的要剪裁的對(duì)象是已訓(xùn)練好的模型，而不是從網(wǎng)絡(luò)結(jié)構(gòu)層進(jìn)行融合處理，最簡(jiǎn)單的方法就是調(diào)整 “-pth(閾值)” ，在精確度與模型大小之間取得平衡，調(diào)整的值越高就會(huì)得到越小的模型，但精準(zhǔn)度會(huì)受到比較大的損失，這就需要經(jīng)過(guò)不斷嘗試去調(diào)節(jié)出滿足要求的模型。

以這個(gè) classification 項(xiàng)目為例，在 80 回合中訓(xùn)練的模型大小為 87.7MB，通常最后一回合訓(xùn)練的模型精確度會(huì)最好，因此用 resnet_080.tlt 模型進(jìn)行修剪測(cè)試：

-pth=0.6：修剪后生成大小為38.5MB的resnet18_nopool_bn_pruned.tlt

-pth=0.3：修剪后生成大小為44.2MB的resnet18_nopool_bn_pruned.tlt

至于修剪后模型對(duì)精確度產(chǎn)生怎樣的影響?就需繼續(xù)往下執(zhí)行“再訓(xùn)練”之后，看看評(píng)估的效果如何!

模型再訓(xùn)練

經(jīng)過(guò)修剪的模型并不能立即使用，而是作為“再訓(xùn)練”的預(yù)訓(xùn)練模型，再次利用遷移學(xué)習(xí)的技巧，以前面的數(shù)據(jù)集進(jìn)行模型再訓(xùn)練的工作，唯一不一樣的地方是這次使用的配置文件為 classification_retrain_spec.cfg，這個(gè)配置內(nèi)容與前面的 classification_spec.cfg 只有三個(gè)差異的地方：

1. 在 model_config 設(shè)置組少了兩個(gè) “freeze_blocks” 設(shè)置;

2. 在 train_config 設(shè)置組的 “pretrained_model_path” 換成前面修剪過(guò)的模型;

3. 在 eval_config 設(shè)置組的 “model_path” 換成修剪后重新訓(xùn)練的模型

繼續(xù)往下執(zhí)行 “7. Retrain pruned models” 與 “8. Testing the model” 兩個(gè)步驟，就能得到修剪過(guò)模型的精確度評(píng)估結(jié)果。

從上面的比較表中可以看出，這個(gè)剪裁的方式對(duì)模型文件大小的影響是很明顯的，但精確度的損失并不是太大，所以這個(gè)剪裁應(yīng)視為有效果。

本實(shí)驗(yàn)結(jié)語(yǔ)

這個(gè)范例主要讓大家體會(huì)并熟悉 TAO 的配置文件與執(zhí)行流程，事實(shí)上 TAO 所提供的范例幾乎都是一樣的流程，后面還有導(dǎo)出模型、INT8 優(yōu)化、生成 TensorRT 引擎的步驟，我們結(jié)合在后面的物件檢測(cè)范例中進(jìn)行說(shuō)明。

原文標(biāo)題：NVIDIA Jetson Nano 2GB 系列文章(60)：圖像分類的模型訓(xùn)練與修剪

文章出處：【微信公眾號(hào)：NVIDIA英偉達(dá)企業(yè)解決方案】歡迎添加關(guān)注!文章轉(zhuǎn)載請(qǐng)注明出處。

審核編輯：湯梓紅