效率對于工業(yè)制造來說至關(guān)重要，即便是微小的提升也會帶來顯著的經(jīng)濟(jì)效益。美國質(zhì)量協(xié)會的數(shù)據(jù)顯示，“許多企業(yè)機(jī)構(gòu)實(shí)際的質(zhì)量相關(guān)成本高達(dá)銷售收入的 15 - 20%，有些甚至高達(dá)總運(yùn)營成本的 40%?！边@些驚人的統(tǒng)計(jì)數(shù)據(jù)揭示了一個(gè)嚴(yán)峻的現(xiàn)實(shí)：工業(yè)應(yīng)用中的缺陷不僅會危及產(chǎn)品質(zhì)量，還會消耗很大一部分的企業(yè)收入。

但如果企業(yè)能夠收回這些損失的利潤，并將其重新投入到創(chuàng)新和擴(kuò)張中，那又會怎樣呢？這正是 AI 的“用武之地”。

本文將探討如何使用NVIDIA TAO設(shè)計(jì)出能夠精準(zhǔn)定位工業(yè)應(yīng)用缺陷，進(jìn)而提高整體質(zhì)量的自定義 AI 模型。

NVIDIA TAO 套件是基于TensorFlow和 PyTorch構(gòu)建的低代碼 AI 工具套件。它通過抽象化 AI 模型和深度學(xué)習(xí)框架的復(fù)雜性，來簡化并加速模型訓(xùn)練過程。借助 TAO 套件，開發(fā)者不僅可以使用預(yù)訓(xùn)練模型，還可以針對特定用例對其進(jìn)行微調(diào)。

在本文中，我們使用一個(gè)名為 VisualChangeNet 的高級預(yù)訓(xùn)練模型進(jìn)行變化檢測，并使用 TAO 工具套件對其進(jìn)行微調(diào)，使其能夠檢測MVTecAnomaly異常檢測數(shù)據(jù)集中的缺陷。這個(gè)綜合基準(zhǔn)數(shù)據(jù)集由各種工業(yè)產(chǎn)品的正常和缺陷樣本組成，專門用于機(jī)器視覺中的異常檢測。

借助 TAO 套件，我們使用遷移學(xué)習(xí)訓(xùn)練了一個(gè)模型，該模型在MVTecAnomaly異常數(shù)據(jù)集的“瓶子”類別上達(dá)到了 99.67% 的總體準(zhǔn)確率、92.3% 的 mIoU、95.8% 的 mF1、97.5 的 mPrecision 和 94.3% 的 mRecall。圖 1 顯示了使用這個(gè)訓(xùn)練有素的模型進(jìn)行的缺陷掩碼預(yù)測。

圖 1. 通過將缺陷圖像與黃金圖像進(jìn)行比較

來分段預(yù)測缺陷物體的缺陷掩碼

第 1 步：設(shè)置先決條件

為了按照本文重新創(chuàng)建這些步驟，請執(zhí)行以下操作。

1. 按照《NGC 用戶指南》中提供的步驟在 NGC 目錄上注冊賬戶并生成 API 密鑰：https://docs.nvidia.com/ngc/gpu-cloud/ngc-user-guide/index.html#generating-api-key

2. 按照《TAO 快速入門指南》設(shè)置 TAO 啟動(dòng)器。為 MVTec 數(shù)據(jù)集下載 VisualChangeNet Segmentation Jupyter Notebook。啟動(dòng) Jupyter Notebook 并運(yùn)行單元格，以便跟著本文進(jìn)行操作。

   *請注意：VisualChangeNet 模型只能在 5.1 及以上的版本中運(yùn)行。

3. 按照提示進(jìn)入下載頁面，下載并準(zhǔn)備MVTec 異常檢測數(shù)據(jù)集，然后復(fù)制 15 個(gè)對象類中任何一個(gè)的下載鏈接：https://www.mvtec.com/company/research/datasets/mvtec-ad

4. 將下載鏈接粘貼到 Jupyter Notebook 第 2.1 部分中的“FIXME”位置，然后運(yùn)行該筆記本單元。本文主要介紹瓶子對象，但所有 15 個(gè)對象都能在該筆記本中使用。圖 2 顯示了數(shù)據(jù)集中的缺陷圖像樣本。

#Download the data
import os
MVTEC_AD_OBJECT_DOWNLOAD_URL = "FIXME"
mvtec_object = MVTEC_AD_OBJECT_DOWNLOAD_URL.split("/")[-1].split(".")[0]
os.environ["URL_DATASET"]=MVTEC_AD_OBJECT_DOWNLOAD_URL
os.environ["MVTEC_OBJECT"]=mvtec_object
!if[!-f$HOST_DATA_DIR/$MVTEC_OBJECT.tar.xz];thenwget$URL_DATASET-O$HOST_DATA_DIR/$MVTEC_OBJECT.tar.xz;elseecho"imagearchivealreadydownloaded";fi

圖 2. MVTec數(shù)據(jù)集中電纜、瓶子

和晶體管的缺陷圖像樣本（從左到右）

我們用 MVTec-AD 中的“瓶子”類別展示了如何在工業(yè)檢測用例中使用 TAO 套件與 VisualChangeNet 來實(shí)現(xiàn)自動(dòng)光學(xué)檢測。

在 Jupyter Notebook 下載數(shù)據(jù)集后，運(yùn)行該筆記的第 2.3 節(jié)，將數(shù)據(jù)集處理成 VisualChangeNet 分割所需的正確格式。

import random 
import shutil 
from PIL import Image
os.environ["HOST_DATA_DIR"] = os.path.join(os.environ["LOCAL_PROJECT_DIR"], "data", "changenet")
formatted_dir = f"formatted_{mvtec_object}_dataset"
 
DATA_DIR = os.environ["HOST_DATA_DIR"]
os.environ["FORMATTED_DATA_DIR"] = formatted_dir
 
#setup dataset folders in expected format 
formatted_path = os.path.join(DATA_DIR, formatted_dir)
a_dir = os.path.join(formatted_path, "A")
b_dir = os.path.join(formatted_path, "B")
label_dir = os.path.join(formatted_path, "label")
list_dir = os.path.join(formatted_path, "list")
 
#Create the expected folders
os.makedirs(formatted_path, exist_ok=True)
os.makedirs(a_dir, exist_ok=True)
os.makedirs(b_dir, exist_ok=True)
os.makedirs(label_dir, exist_ok=True)
os.makedirs(list_dir,exist_ok=True)

該原始數(shù)據(jù)集專門用于異常檢測。我們將兩者合并成一個(gè)包含 283 幅圖像的綜合數(shù)據(jù)集，然后將其分成 253 幅訓(xùn)練集圖像和 30 幅測試集圖像。兩個(gè)數(shù)據(jù)集都包含有缺陷的樣本。

我們確保該測試集包含每個(gè)缺陷類中 30% 的缺陷樣本，“瓶子”類別主要包含“無缺陷”圖像，三個(gè)缺陷類別中，每個(gè)類別各有約 20 幅圖像。

圖 3. 該數(shù)據(jù)集中的輸入樣本，包含測試圖像、黃金圖像和顯示缺陷的分割掩碼。圖中采用從瓶子頂部俯視的視角，攝像頭安裝在瓶口的正上方。

第 2 步：下載 VisualChangeNet 模型

VisualChangeNet 模型是最先進(jìn)的基于轉(zhuǎn)換器的變化檢測模型。其設(shè)計(jì)核心是孿生網(wǎng)絡(luò)（Siamese Network）。孿生網(wǎng)絡(luò)是一種獨(dú)特的神經(jīng)網(wǎng)絡(luò)架構(gòu)，由兩個(gè)或多個(gè)相同的子網(wǎng)絡(luò)組成。這些“孿生”子網(wǎng)絡(luò)接受不同的輸入，但共享相同的參數(shù)和權(quán)重。這種架構(gòu)使 VisualChangeNet 模型能夠比較當(dāng)前圖像和作為參考的“黃金”圖像之間的特征，從而精準(zhǔn)定位變化和變更。該功能使孿生網(wǎng)絡(luò)尤其擅長圖像對比和異常檢測等任務(wù)。

模型文檔提供更多細(xì)節(jié)，如架構(gòu)和訓(xùn)練數(shù)據(jù)等。我們沒有從頭開始訓(xùn)練模型，而是從在 NV-ImageNet 數(shù)據(jù)集上訓(xùn)練而成的預(yù)訓(xùn)練 FAN 骨干開始。我們使用 TAO 套件在 MVTec-AD 數(shù)據(jù)集上針對“瓶子”類別對其進(jìn)行微調(diào)。

運(yùn)行筆記第 3 部分以安裝 NGC 命令行工具，并從 NGC 下載該預(yù)訓(xùn)練骨干。

# Installing NGC CLI on the local machine.
## Download and install
import os
%env CLI=ngccli_cat_linux.zip
!mkdir -p $HOST_RESULTS_DIR/ngccli
 
# # Remove any previously existing CLI installations
!rm -rf $HOST_RESULTS_DIR/ngccli/*
!wget "https://ngc.nvidia.com/downloads/$CLI" -P $HOST_RESULTS_DIR/ngccli
!unzip -u "$HOST_RESULTS_DIR/ngccli/$CLI" -d $HOST_RESULTS_DIR/ngccli/
!rm $HOST_RESULTS_DIR/ngccli/*.zip
os.environ["PATH"]="{}/ngccli/ngc-cli:{}".format(os.getenv("HOST_RESULTS_DIR", ""), os.getenv("PATH", ""))
!mkdir -p $HOST_RESULTS_DIR/pretrained
!ngc registry model list nvidia/tao/pretrained_fan_classification_nvimagenet*
!ngcregistrymodeldownload-version"nvidia/tao/pretrained_fan_classification_nvimagenet:fan_base_hybrid_nvimagenet"--dest$HOST_RESULTS_DIR/pretrained

第 3 步：使用 TAO 套件訓(xùn)練模型

在本節(jié)中，我們將詳細(xì)介紹如何使用 TAO 套件訓(xùn)練 VisualChangeNet 模型。您可在模型卡中找到 VisualChangeNet 模型的詳細(xì)信息以及支持的預(yù)訓(xùn)練權(quán)重，也可以使用該預(yù)訓(xùn)練 FAN 主干權(quán)重作為微調(diào) VisualChangeNet 的起點(diǎn)，這也是我們在 MVTec-AD 數(shù)據(jù)集上進(jìn)行微調(diào)時(shí)使用的權(quán)重。

如圖 4 所示，該訓(xùn)練算法會同步更新所有子網(wǎng)絡(luò)的參數(shù)。在 TAO 中，VisualChangeNet 支持將兩幅圖像作為黃金樣本和測試樣本輸入，其目標(biāo)是檢測“黃金或參考”圖像與“測試”圖像之間的變化。TAO 支持 VisualChangeNet 架構(gòu)的 FAN 骨干網(wǎng)絡(luò)。

TAO 支持兩類變化檢測網(wǎng)絡(luò)，分別是 VisualChangeNet-Segmentation 和 VisualChangeNet-Classification。在本文中，我們將使用 VisualChangeNet-Segmentation 模型，通過分割 MVTec-AD 數(shù)據(jù)集中兩張輸入圖像之間的變化像素，來演示變化檢測。

圖 4. 檢測瓶子類黃金圖像和測試圖像之間變化的

VisualChangeNet-Segmentation 算法架構(gòu)圖

使用工具套件微調(diào) VisualChangeNet 模型非常簡單，無需編碼經(jīng)驗(yàn)。只需在 TAO 套件中加載數(shù)據(jù)、設(shè)置實(shí)驗(yàn)配置并運(yùn)行訓(xùn)練命令即可。

該實(shí)驗(yàn)配置文件定義了 VisualChangeNet 模型架構(gòu)、訓(xùn)練和評估的超參數(shù)。可在訓(xùn)練模型前在 Jupyter Notebook 中查看并編輯該配置文件。

我們使用該配置對 VisualChangeNet 模型進(jìn)行微調(diào)，在該配置中定義了一個(gè)帶有預(yù)訓(xùn)練 FAN-Hybrid-Base 骨干（即基準(zhǔn)模型）的 VisualChangeNet 模型，然后對該模型進(jìn)行 30 個(gè)歷元訓(xùn)練，批量大小為 8。下面展示了部分實(shí)驗(yàn)配置以及一些關(guān)鍵參數(shù)。完整的實(shí)驗(yàn)配置可在 Jupyter Notebook 中進(jìn)行查看。

encryption_key: tlt_encode
task: segment
train:
  resume_training_checkpoint_path: null
  pretrained_model_path: null
  segment:
    loss: "ce"
    weights: [0.5, 0.5, 0.5, 0.8, 1.0]
  num_epochs: 30
  num_nodes: 1
  val_interval: 1
  checkpoint_interval: 1
  optim:
    lr: 0.0002
    optim: "adamw"
    policy: "linear"
    momentum: 0.9
    weight_decay: 0.01
results_dir: "/results"
model:
  backbone:
    type: "fan_base_16_p4_hybrid"
pretrained_backbone_path:/results/pretrained/pretrained_fan_classification_nvimagenet_vfan_base_hybrid_nvimagenet/fan_base_hybrid_nvimagenet.pth

為了調(diào)整模型的性能，可以修改的一些常用值包括訓(xùn)練歷元數(shù)、學(xué)習(xí)率（lr）、優(yōu)化器和預(yù)訓(xùn)練骨干。如要從頭開始訓(xùn)練，可以將 pretrained_backbone_path 設(shè)置為空，但這可能會增加為達(dá)到高準(zhǔn)確率所需的歷元數(shù)和數(shù)據(jù)量。如需進(jìn)一步了解該實(shí)驗(yàn)配置文件中的參數(shù)，請參閱《VisualChangeNet用戶指南》：https://docs.nvidia.com/tao/tao-toolkit/text/visual_changenet/index.html

現(xiàn)在數(shù)據(jù)集和實(shí)驗(yàn)配置已經(jīng)準(zhǔn)備就緒，我們要開始在 TAO 套件中進(jìn)行訓(xùn)練了。運(yùn)行第 5.1 節(jié)中的代碼塊，使用單顆 GPU 啟動(dòng) VisualChangeNet 訓(xùn)練。

print("Train model")
!tao model visual_changenet train 
                  -e $SPECS_DIR/experiment.yaml 
                    train.num_epochs=$NUM_EPOCHS 
                    dataset.segment.root_dir=$DATA_DIR 
model.backbone.pretrained_backbone_path=$BACKBONE_PATH

該單元將開始在 MVTec 數(shù)據(jù)集上訓(xùn)練 VisualChangeNet-Segmentation 模型。在訓(xùn)練過程中，模型將學(xué)習(xí)如何識別缺陷對象并輸出顯示缺陷區(qū)域的分割掩碼。訓(xùn)練日志包含驗(yàn)證數(shù)據(jù)集上的準(zhǔn)確率、訓(xùn)練損失、學(xué)習(xí)率和訓(xùn)練后的模型等，都將保存在實(shí)驗(yàn)配置中設(shè)置的結(jié)果目錄中。

第 4 步：評估模型

訓(xùn)練完成后，我們可以使用 TAO 在驗(yàn)證數(shù)據(jù)集上對模型進(jìn)行評估。VisualChangeNet Segmentation 的輸出是 2 張給定輸入圖像的分割變化圖，表示的是像素級缺陷。筆記的第 6 節(jié)將運(yùn)行評估模型性能的命令。

!tao model visual_changenet evaluate 
                   -e $SPECS_DIR/experiment.yaml 
                    evaluate.checkpoint=$RESULTS_DIR/train/changenet.pth 
                    dataset.segment.root_dir=$DATA_DIR

TAO 中的評估命令將返回驗(yàn)證集的多個(gè) KPI，比如在缺陷類（缺陷像素）上的準(zhǔn)確度、精確度、召回率、F1 分?jǐn)?shù)和 IoU 等。

OA = 有變化/無變化像素的總體準(zhǔn)確率（輸入維度 - 256×256）

表 1. MVTec-AD 二進(jìn)制變化檢測（瓶子類）

中的 VisualChangeNet 模型評估指標(biāo)

第 5 步：部署模型

可以使用NVIDIA DeepStream或NVIDIA Triton部署這一經(jīng)過微調(diào)的模型。我們將其導(dǎo)出為 .onnx 格式。筆記第 8 節(jié)將運(yùn)行 TAO 導(dǎo)出命令。

!tao model visual_changenet export 
                    -e $SPECS_DIR/experiment.yaml 
                        export.checkpoint=$RESULTS_DIR/train/changenet.pth 
                        export.onnx_file=$RESULTS_DIR/export/changenet.onnx

輸出的 .onnx 模型與訓(xùn)練好的 .pth 模型保存在同一目錄下。如要部署到 Triton，請查看 GitHub 上的 tao-toolkit-triton 資源庫。該項(xiàng)目提供了將許多 TAO 模型（包括 VisualChangeNet Segmentation）部署到 Triton 推理服務(wù)器的參考實(shí)現(xiàn)。

實(shí)時(shí)推理性能

推理以 FP16 精度在所提供的未修剪模型上運(yùn)行。該推理性能在嵌入式Jetson OrinGPU 和數(shù)據(jù)中心 GPU 上使用 trtexec 運(yùn)行。Jetson 設(shè)備以 GPU 頻率最高的 Max-N 配置運(yùn)行。

請執(zhí)行以下命令來運(yùn)行 trtexec：

/usr/src/tensorrt/bin/trtexec --onnx=<ONNX path> --minShapes=input0:1x3x512x512,input1:1x3x512x512 --maxShapes=input0:8x3x512x512,input1:8x3x512x512 --optShapes=input0:4x3x512x512,input1:4x3x512x512 
--saveEngine=<engine path>

此處顯示的性能僅為推理性能。流式傳輸視頻數(shù)據(jù)的端到端性能可能會因硬件和軟件的其他瓶頸而有所變化。

總結(jié)

在本文中，我們了解了如何使用 TAO 套件微調(diào) VisualChangeNet 模型，并將其用于分割 MVTec 數(shù)據(jù)集中的缺陷，從而實(shí)現(xiàn) 99.67% 的總體準(zhǔn)確率。

現(xiàn)在，您還可以使用 NVIDIA TAO 檢測制造工作流程中的缺陷。

如要開始使用，請先：

• 從 NVIDIA NGC 目錄下載 VisualChangeNet 模型。

• 按照《TAO 快速入門指南》設(shè)置 TAO 啟動(dòng)器：https://docs.nvidia.com/tao/tao-toolkit/text/tao_toolkit_quick_start_guide.html

• 從GitHub下載 VisualChangeNet Segmentation Notebook：https://github.com/NVIDIA/tao_tutorials/tree/main/notebooks/tao_launcher_starter_kit/visual_changenet

• 請通過NVIDIA 文檔進(jìn)一步了解 NVIDIA TAO 工具套件：https://docs.nvidia.com/tao/tao-toolkit/text/overview.html#:~:text=NVIDIA%20TAO%20Toolkit%20is%20a,and%20the%20deep%20learning%20framework

GTC 2024 將于 2024 年 3 月 18 至 21 日在美國加州圣何塞會議中心舉行，線上大會也將同期開放。點(diǎn)擊“閱讀原文”或掃描下方海報(bào)二維碼，立即注冊 GTC 大會。