TAO是英伟达推出的可以简化企业AI应用和部署的模型自适应平台,其提供了大量的预训练模型以及模型finetune程序,使得用户可以快速的在自己的数据上进行模型finetune,之后用TAO的部署工具可以快速完成模型的部署。
TAO提供了四种使用方式,本文将介绍其中的容器使用方法。官方手册:Working With the Containers — TAO Toolkit 4.0 documentationhttps://docs.nvidia.com/tao/tao-toolkit/text/working_with_the_containers.html
要使用TAO的镜像首先需要注册NV-NGC账号,之后使用 docker login nvcr.io登陆账号:
a. Username: "$oauthtoken"
b. Password: "YOUR_NGC_API_KEY"
其中Username统一使用$oauthtoken,Password使用NV-NGC登陆账户的set up中生成的API key,具体参考下图:
点击右上角的按键生成如下的key:
复制上图生成的key登陆出现successed即可。
下面介绍TAO预训练模型的使用方法,本文用TAO提供的车牌号识别的预训练模型为样例。
LPRNet预训练模型-车牌号识别
官方手册:https://docs.nvidia.com/tao/tao-toolkit/text/character_recognition/lprnet.html#
模型介绍:License Plate Recognition | NVIDIA NGC
TAO提供了好几个镜像,里面分别包含不同的模型,如下图:
根据需要使用的预训练模型查找相应的镜像,可参考容器手册:Working With the Containers — TAO Toolkit 4.0 documentationhttps://docs.nvidia.com/tao/tao-toolkit/text/working_with_the_containers.html
如下图所示,LRNNet包含在tao-toolkit:4.0.0-tf1.15.5,
因此直接在docker里pull这个镜像即可,镜像地址:TAO Toolkit | NVIDIA NGCDocker containers distributed as part of the TAO Toolkit packagehttps://catalog.ngc.nvidia.com/orgs/nvidia/teams/tao/containers/tao-toolkit/tags
下载LRNNet的预训练模型,地址:License Plate Recognition | NVIDIA NGCModel to recognize characters from the image crop of a License Plate.https://catalog.ngc.nvidia.com/orgs/nvidia/teams/tao/models/lprnet/version,其中分别为一个可训练的模型和可部署的模型,根据需要下载即可。
完成预训练模型的下载和准备好镜像镜像之后启动容器,记得将模型映射到容器里,此处不在赘述容器的启动方法。
预训练模型推理
TAO将LPR的推理、训练、验证等代码封装打包成了一个程序,因此进入容器之后可直接使用lprnet程序来完成相应的功能,用lprnet -h可查看相关说明,如下图:
可以看到lprnet支持验证、导出模型、推理和训练四个task,导出模型是将用于训练的tlt模型导出为etlt模型以便于进行部署。官方手册运行推理的命令地址:LPRNet — TAO Toolkit 4.0 documentation
https://docs.nvidia.com/tao/tao-toolkit/text/character_recognition/lprnet.html#running-inference-on-the-lprnet-model
由推理的命令可知我们需要提供测试图像、模型位置和配置文件,配置文件的编写在官方手册有介绍,地址:LPRNet — TAO Toolkit 4.0 documentationhttps://docs.nvidia.com/tao/tao-toolkit/text/character_recognition/lprnet.html#creating-an-experiment-spec-file
值得注意的是配置文件中字符集文件的参数要填写下载的预训练模型包含的字符集路径,如下图
下载预训练模型时另外另个文件就是字符集,上图红框中的参数要写下图红框文件的保存路径。根据识别中国或美国车牌的需求选择一个即可。配置文件要保存成.prototxt格式!!!
本文随机在网上下载了一张中文车牌号图片进行测试,如下图:
最后使用命令:lprnet inference --gpu_index=0 -m 模型路径 -i 数据文件夹路径 -e 配置文件路径 -k 模型解码秘钥。
需要注意的是目前TAO的预训练模型都进行了加密,因此在使用模型的时候需要提供解码的秘钥,秘钥可以在模型的介绍页面进行查找,本文使用的车牌识别模型在介绍页面提供了解码秘钥:Model load key: nvidia_tlt。目前TAO所有预训练模型的key应该都是nvidia_tlt。
本文运行结果如下图:红框中就是识别的结果。
预训练模型Finetune
如果要用自己的数据对TAO提供的预训练模型进行finetune/train,只需在配置文件中修改数据的路径及字符集路径即可,
训练数据的准备参考:LPRNet — TAO Toolkit 4.0 documentationhttps://docs.nvidia.com/tao/tao-toolkit/text/character_recognition/lprnet.html#preparing-the-dataset
最后训练命令参考:
LPRNet — TAO Toolkit 4.0 documentationhttps://docs.nvidia.com/tao/tao-toolkit/text/character_recognition/lprnet.html#training-the-model部署
模型部署可参考:LPRNet — TAO Toolkit 4.0 documentationhttps://docs.nvidia.com/tao/tao-toolkit/text/character_recognition/lprnet.html#deploying-the-model
部署需要使用到TAO的deploy镜像,地址:TAO Toolkit | NVIDIA NGCDocker containers distributed as part of the TAO Toolkit packagehttps://catalog.ngc.nvidia.com/orgs/nvidia/teams/tao/containers/tao-toolkit/tags
部署需要使用etlt模型,因此训练好的模型需要导出成etlt格式,然后按照官方手册用tao-converter将etlt模型转为TensorRT的engine,之后使用Deep streamSDK进行部署即可,Deep streamSDK的用法不在本文介绍。