效率在工業制造中至關重要，在工業制造中，即使是微小的收益也會產生重大的財務影響。據美國質量協會稱，“許多組織的真正質量相關成本將高達銷售收入的 15-20%，有些則高達總運營的 40%.”這些驚人的統計數據揭示了一個嚴峻的現實：工業應用中的缺陷不僅會損害產品質量，而且會消耗公司收入的很大一部分。
但是，如果公司能夠收回這些損失的利潤，并將其重新用于創新和擴展呢？這是 AI 潛力的亮點所在。
本文將探討如何使用 NVIDIA TAO 設計自定義 AI 模型，以找出工業應用中的缺陷，從而提高整體質量。
NVIDIA TAO 工具套件是基于 TensorFlow 和 PyTorch 構建的低代碼 AI 工具包。它通過抽象出 AI 模型和深度學習框架的復雜性來簡化和加速模型訓練過程。借助 TAO 工具套件，開發者可以使用預訓練模型，并針對特定用例對其進行微調。
在本文中，我們利用名為 VisualChangeNet 的高級預訓練模型進行變化檢測，并使用 TAO 工具套件對其進行微調，以檢測 MV Tech Anomaly 檢測數據集中的缺陷。這個全面的基準測試數據集專為機器視覺中的異常檢測而設計，由包含正常和有缺陷樣本的各種工業產品組成。
借助 TAO 工具套件，我們使用遷移學習訓練模型，該模型在 MVTec Anomaly 數據集的瓶子類上的總體準確度分別為 99.67%、92.3%mIoU、95.8%mF1、97.5 mPrecision 和 94.3%mRecall.圖 1 顯示了使用經過訓練的模型進行的缺陷掩膜預測。
第 1 步：安裝前提條件
要遵循帖子并重新創建這些步驟，請執行以下操作。
- 在 NGC 目錄中注冊帳戶，并按照 NGC 提供的步驟生成 API 密鑰，詳情請參閱用戶指南。
- 按照 TAO 快速入門指南設置 TAO 啟動器。下載適用于 MVTec 數據集的 VisualChangeNet Segmentation Jupyter Notebook。啟動 Jupyter Notebook 并運行此博文中關注的單元。
*請注意，VisualChangeNet 模型僅適用于 5.1 版本。 - 下載并準備 MVTec 異常檢測數據集，只需按照下載頁面上的提示操作，并復制任意一個 15 個對象類別的下載鏈接即可。
- 將下載鏈接粘貼到 Jupyter Notebook 第 2.1 節中的“FIXME”位置，并運行 Notebook 單元。本文重點介紹瓶子對象，但所有 15 個對象都在 Notebook 中運行。圖 2 顯示了數據集中的樣本缺陷圖像。
#Download the dataimport osMVTEC_AD_OBJECT_DOWNLOAD_URL = "FIXME"mvtec_object = MVTEC_AD_OBJECT_DOWNLOAD_URL.split("/")[-1].split(".")[0]os.environ["URL_DATASET"]=MVTEC_AD_OBJECT_DOWNLOAD_URLos.environ["MVTEC_OBJECT"]=mvtec_object!if [ ! -f $HOST_DATA_DIR/$MVTEC_OBJECT.tar.xz ]; then wget $URL_DATASET -O $HOST_DATA_DIR/$MVTEC_OBJECT.tar.xz; else echo "image archive already downloaded"; fi |
在 MVTec-AD 中，我們利用 TAO 工具套件，利用 VisualChangeNet 展示工業檢測用例的自動光學檢測。
Jupyter Notebook 下載數據集后，運行 Notebook 第 2.3 節，將數據集處理為正確的格式，以進行 VisualChangeNet 分割。
import random import shutil from PIL import Imageos.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 foldersos.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) |
原始數據集專為異常檢測而設計。我們將兩者合并以創建包含 283 張圖像的組合數據集，然后將其分為 253 張訓練集圖像和 30 張測試集圖像。這兩套圖像都包含有缺陷的樣本。
我們確保測試集包含每個缺陷類別中 30%的缺陷樣本，因為`bottle`類主要包含`no-defect`圖像，三個缺陷類別中每個類別大約有 20 張圖像。
第 2 步：下載 VisualChangeNet 模型
VisualChangeNet 模型是基于 Transformer 的先進變化檢測模型。其設計的核心是 Siamese 網絡。Siamese 網絡是由兩個或多個相同的子網絡組成的獨特的神經網絡架構。這些“生”子網絡接受不同的輸入，但共享相同的參數和權重。在 VisualChangeNet 的上下文中，此架構使模型能夠比較當前圖像和參考“黃金”圖像之間的特征，精確定位變化和變化。這種能力使 Siamese Networks 特別擅長處理圖像比較和異常檢測等任務。
我們的模型文檔提供了更多細節，例如模型架構和訓練數據。我們不是從零開始訓練模型，而是首先利用在 NV-ImageNet 數據集上進行訓練的預訓練的 FAN 主干網絡。然后，我們使用 TAO 工具套件在 MVTec-AD 數據集上對 Bottle 類進行微調。
運行 Notebook 的第 3 節以安裝 NGC 命令行工具，并從 NGC 下載預訓練的骨干網絡。
# Installing NGC CLI on the local machine.## Download and installimport 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/*!unzip -u "$HOST_RESULTS_DIR/ngccli/$CLI" -d $HOST_RESULTS_DIR/ngccli/!rm $HOST_RESULTS_DIR/ngccli/*.zipos.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*!ngc registry model download-version "nvidia/tao/pretrained_fan_classification_nvimagenet:fan_base_hybrid_nvimagenet" --dest $HOST_RESULTS_DIR/pretrained |
第 3 步：使用 TAO 工具套件訓練模型
在本節中，我們將詳細介紹如何使用 TAO 工具套件訓練 VisualChangeNet 模型。您可以在工具套件中找到 Visual ChangeNet 模型的詳細信息以及受支持的預訓練權重 模型卡。您還可以使用預訓練的 FAN 主干權重作為微調 VisualChangeNet 的起點，這是我們在 MVTec-AD 數據集上進行微調的起點。
如圖 4 所示，訓練算法會同時更新所有子網絡中的參數。在 TAO 中，Visual ChangeNet 支持將兩個圖像作為輸入 – 一個黃金樣本和一個測試樣本。目標是檢測“黃金或參考”圖像與“測試”圖像之間的變化。TAO 支持 FAN 作為可視化 ChangeNet 架構的骨干網絡。
TAO 支持兩種類型的變化檢測網絡：Visual ChangeNet-Segmentation 和 Visual ChangeNet-Classification.在本文中，我們利用 Visual ChangeNet – Segmentation 模型，通過在 MVTec-AD 數據集中的兩個輸入圖像之間分割更改的像素來演示變化檢測。
使用 TAO 工具套件可以輕松微調 VisualChangeNet 模型，并且無需任何編碼經驗。只需在 TAO 工具套件中加載數據，設置實驗配置，然后運行 train 命令。
實驗配置文件定義了 VisualChangeNet 模型架構、訓練和評估的超參數。在 Jupyter Notebook 中，您可以在訓練模型之前查看和編輯配置文件。
我們使用此配置來微調 Visual ChangeNet 模型。在配置中，讓我們使用預訓練的 FAN-Hybrid-Base 主干(即基準模型)定義一個 Visual ChangeNet 模型。讓我們使用批量大小為 8 的 30 次迭代來訓練模型。下一節演示了部分實驗配置，其中顯示了一些關鍵參數。完整的實驗配置可在 Jupyter Notebook 中查看。
encryption_key: tlt_encodetask: segmenttrain: 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.01results_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 |
您可以修改一些常用參數以調整模型的性能，包括訓練次數、學習率 (lr)、優化器和預訓練的主干網絡。若要從頭開始訓練，可以將 pretrained_backbone_path 設置為 null，但這可能會增加所需的訓練次數和實現高準確度所需的數據量。有關實驗配置文件中參數的更多信息，請參閱VisualChangeNet 用戶指南。
現在，數據集和實驗配置已準備就緒，我們開始在 TAO 工具套件中進行訓練。運行第 5.1 節中的代碼塊，使用單個 GPU 啟動 Visual ChangeNet 訓練。
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 數據集上訓練 Visual ChangeNet Segmentation 模型。在訓練期間，該模型將學習如何識別缺陷對象并輸出顯示缺陷區域的分割掩膜。訓練日志(包括驗證數據集的準確性、訓練損失、學習率和經過訓練的模型)保存在實驗配置中設置的結果目錄中。
第 4 步：評估模型
訓練完成后，我們可以使用 TAO 評估驗證數據集上的模型。對于 Visual ChangeNet Segmentation，輸出是 2 個給定輸入圖像的分割更改圖，表示像素級缺陷。Notebook 的第 6 節將運行該命令來評估模型的性能。
!tao model visual_changenet evaluate \ -e $SPECS_DIR/experiment.yaml \ evaluate.checkpoint=$RESULTS_DIR/train/changenet.pth \ dataset.segment.root_dir=$DATA_DIR |
TAO 中的 evaluate 命令將返回驗證集上的多個 KPI，例如準確性、精度、召回率、F1 分數和缺陷類(缺陷像素)的 IoU.
OA=更改/不更改像素的總體準確性(輸入維度 – 256*256)
MVTec-AD 二進制 CD（Bottle 類）
| 模型 | 骨干網絡 | 米精度 | mRecall | mF1 | mIOU | OA |
| VisualChangeNet | FAN-Hybrid-B (預訓練) | 97.5% | 94.3% | 95.8 | 92.3 | 99.67％ |
第 5 步：部署模型
您可以使用此微調模型，并利用 NVIDIA DeepStream 或 NVIDIA Triton 將其導出為 .onnx 格式。notebook 的第 8 節將介紹如何運行 TAO 導出命令。
!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 模型與經過訓練的 .pth 模型保存在同一目錄中。要部署到 Triton，請查看 tao-toolkit-triton GitHub 上的存儲庫。此項目提供參考實現，以將許多 TAO 模型（包括 Visual ChangeNet Segmentation）部署到 Triton 推理服務器。
實時推理性能
在提供的未刪減模型上以 FP16 精度運行推理。在嵌入式 Jetson Orin GPU 和數據中心 GPU 上使用 trtexec 運行推理性能。Jetson 設備以 Max-N 配置運行，以實現最大 GPU 頻率。
運行以下命令以運行 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>
此處顯示的性能是純推理性能。流式傳輸視頻數據的端到端性能可能會因硬件和軟件中的其他瓶頸而有所不同。
| 平臺 | 批量大小 | FPS |
|---|---|---|
| NVIDIA Jetson Orin Nano 8GB | 16 | 15.19 |
| NVIDIA Jetson Orin NX 16 GB | 16 | 21.92% |
| NVIDIA Jetson AGX Orin 64 GB | 16 | 55.07 |
| NVIDIA A2 Tensor Core GPU | 16 | 36.02 |
| NVIDIA T4 Tensor Core GPU | 16 | 59.7 |
| NVIDIA L4 Tensor Core GPU | 8 | 131.48 |
| NVIDIA A30 Tensor Core GPU | 16 | 204.12 |
| NVIDIA L40 GPU | 8 | 364 |
| NVIDIA A100 Tensor Core GPU | 32 | 435.18 |
| NVIDIA H100 Tensor Core GPU | 32 | 841.68 |
總結
在本文中，我們學習了如何使用 TAO 工具套件微調 VisualChangeNet 模型，并將其用于分割 MVTech 數據集中的缺陷，總體準確率達到 99.67％。
現在，您還可以利用 NVIDIA TAO 檢測制造工作流程中的缺陷。
要開始使用，請執行以下操作：
- 從 NVIDIA NGC 目錄中下載 VisualChangeNet 模型。
- 按照 TAO 快速入門指南 來設置并啟動 TAO 工具包。
- 從以下網址下載 Visual ChangeNet 分割筆記本：GitHub
通過以下鏈接詳細了解 NVIDIA TAO 工具套件：NVIDIA 文檔。
?
