自定义 AI 模型：使用 NVIDIA Triton 部署字符检测和识别模

这篇文章最初发表在 NVIDIA 技术博客上。

NVIDIA Triton Inference Server 通过使团队能够在任何基于 GPU 或 CPU 的基础设施上部署、运行和扩展经过训练的 ML 或 DL 模型，简化和标准化 AI 推理。它帮助开发人员在云端、本地、边缘和嵌入式设备上提供高性能推理。

nvOCDR 库已集成到 Triton 中进行推理。nvOCDR 库封装了用于光学字符检测和识别（OCD/OCR）的整个推理管道。该库使用在 TAO Toolkit 上训练的 OCDNet 和 OCRNet 模型。想要了解更多详细信息，请参阅 nvOCDR 文档。

本文是关于使用 NVIDIA TAO 和预训练模型创建和部署自定义 AI 模型以准确检测和识别手写文本的系列文章的一部分。第一部分解释了如何使用 TAO 对字符检测和识别模型进行训练和微调。本部分将引导您完成使用 NVIDIA Triton 部署模型的步骤。所提供的步骤可用于任何其他 OCR 任务。

使用 OCD/OCR 模型构建 Triton 样本

以下步骤显示了在 Triton 推理服务器中使用 Docker 图像构建和使用 OCD/OCR 模型的简单推荐方法。

步骤 1：准备 ONNX 模型

一旦你按照ocdnet.ipynb和ocrnet.ipynb要完成模型训练和导出，可以获得两个 ONNX 模型，例如ocdnet.onnx和ocrnet.onnx.（在ocdnet.ipynb 里，导出的 ONNX 命名为model_best.onnx在里面ocrnet.ipynb，导出的 ONNX 命名为best_accuracy.onnx）

shell 复制代码

# bash commands
$ mkdir onnx_models
$ cd onnx_models
$ cp <ocd results dir>/export/model_best.onnx ./ocdnet.onnx
$ cp <ocr results dir>/export/best_accuracy.onnx ./ocrnet.onnx

字符列表文件，生成于ocrnet.ipynb，也是必需的：

shell 复制代码

$ cp <ocr DATA_DIR>/character_list ./

步骤 2：获取 nvOCDR 存储库

要获取 nvOCDR 存储库，请使用以下脚本：

shell 复制代码

$ git clone https://github.com/NVIDIA-AI-IOT/NVIDIA-Optical-Character-Detection-and-Recognition-Solution.git

步骤 3：构建 Triton 服务器 Docker 映像

Triton 服务器和客户端 Docker 镜像的构建过程可以通过运行相关脚本自动启动：

shell 复制代码

$ cd NVIDIA-Optical-Character-Detection-and-Recognition-Solution/triton

# bash setup_triton_server.sh [input image height] [input image width] [OCD input max batchsize] [DEVICE] [ocd onnx path] [ocr onnx path] [ocr character list path]
$ bash setup_triton_server.sh 1024 1024 4 0 ~/onnx_models/ocd.onnx ~/onnx_models/ocr.onnx ~/onnx_models/ocr_character_list

步骤 4：构建 Triton 客户端 Docker 镜像

使用以下脚本构建 Triton 客户端 Docker 映像：

shell 复制代码

$ cd NVIDIA-Optical-Character-Detection-and-Recognition-Solution/triton
$ bash setup_triton_client.sh

步骤 5：运行 nvOCDR Triton 服务器

在构建 Triton 服务器和 Triton client Docker 镜像后，创建一个容器并启动 Triton'服务器：

css 复制代码

$ docker run -it --net=host --gpus all --shm-size 8g nvcr.io/nvidian/tao/nvocdr_triton_server:v1.0 bash

接下来，修改 nvOCDR-lib 的配置文件。nvOCDR-lib 可以支持高分辨率输入图像（4000 x 4000 或更大）。如果您的输入图像很大，您可以在 Triton 服务器容器中将配置文件更改为/opt/nvocdr/ocdr/triton/models/nvOCDR/spec.json，以支持高分辨率图像推断。

bash 复制代码

# to support high resolution images
$ vim /opt/nvocdr/ocdr/triton/models/nvOCDR/spec.json

   "is_high_resolution_input": true,
   "resize_keep_aspect_ratio": true,

这个resize_keep_aspect_ratio将自动设置为 Trueis_high_resolution_input为 True。如果您要推断分辨率较小的图像（例如 640 x 640 或 960 x 1280），您可以设置 is_high_resolution_input为 False。

在容器中，运行以下命令以启动 Triton 服务器：

ini 复制代码

$ CUDA_VISIBLE_DEVICES=<gpu idx> tritonserver --model-repository /opt/nvocdr/ocdr/triton/models/

步骤 6：发送推理请求

在一个单独的控制台中，从 Triton 客户端容器启动 nvOCDR 示例：

shell 复制代码

$ docker run -it --rm -v <path to images dir>:<path to images dir>  --net=host nvcr.io/nvidian/tao/nvocdr_triton_client:v1.0 bash

启动推理：

shell 复制代码

$ python3 client.py -d <path to images dir> -bs 1

图 1。OCDNet 和 OCRNet 对样本手写图像的预测输出

结论

NVIDIA TAO 5.0 介绍了光学字符检测（OCD）和光学字符识别（OCR）的几个功能和模型。本文介绍了自定义和微调预训练模型的步骤，以准确识别 IAM 数据集上的手写文本。该模型实现了 90%的字符检测准确率和 80%的字符识别准确率。文章中提到的所有步骤都可以从提供的 Jupyter 笔记本上运行，从而可以轻松地用最少的代码创建自定义的人工智能模型。

有关详细信息，请参阅：

阅读原文