众所周知,深度学习的训练比较玄学,大家经常调侃就像"炼丹"一样。如果你有个好工具,科学"炼丹"的效率就会显著提升!
Amusi 这里给大家介绍的是 NVIDIA 官方推出的 TAO 工具套件,即一个基于 Python 的工具包,通过优化预训练模型和应用迁移学习来加速模型训练以实现更高性能的 AI 系统,目前支持 TensorFlow、PyTorch 等深度学习框架。
TAO 工具套件
NVIDIA 训练、调整和优化 (TAO) 工具套件 (TAO:Train, Adapt, and Optimize)是一款基于 CLI 和 Jupyter Notebook 的低代码(low-code) AI 模型开发解决方案, 方便对 NVIDIA 预训练 AI 模型进行调优,主要支持计算机视觉和对话式 AI 两大领域。
TAO 工具套件特性
- 训练更轻松
- AI 高度准确
- 推理优化(低延迟/高吞吐)
- 部署更容易
目前 TAO 工具套件支持超过 100 种经 NVIDIA 优化的模型架构和 Backbone (主干) 的组合,例如图像分类模型:ResNet、VGG、EfficientNet 等;目标检测模型:YOLOv3/v4、RetinaNet、Faster R-CNN 等;分割模型:Mask R-CNN、U-Net等。
此外,还有一些偏应用落地类模型,比如可用于新零售领域和社交距离计算的 PeopleNet;可用于停车记录的 License Plate Detection 等。
NVIDIA 预训练库提供的部分 CV 模型及 Backbone 如下图所示:
当你选择好合适的模型之后,就可以使用 TAO 工具套件来训练、微调、剪枝和导出高度优化且准确的 AI 模型,最后,还可以用部署框架将生产级模型部署应用起来,比如计算机视觉模型可以搭配 DeepStream,而对话式 AI 模型可以搭配 Riva 来构建端到端服务和部署方案 。
TAO 工具套件流程框架图
NVIDIA PeopleNet 行人检测的演示效果图
TAO 工具套件 3.0-22.05 版本亮点
就在前些天,NVIDIA 发布了最新版本的 TAO 工具套件(版本 3.0-22.05 )。该最新版套件提供了更多的视觉和语音预训练模型,还上新了很多新功能,比如支持了 TensorBoard、自定义 ONNX 模型权重导入等。
值得说一下,从 TAO Toolkit 3.0-22.05 开始,提供的目标检测模型都支持 COCO 格式。下面将重点介绍 TAO 最新版本的部分亮点:
1. 使用 TensorBoard 进行可视化
训练过深度学习模型的同学应该都知道大名鼎鼎的 TensorBoard,可以称为 TensorFlow 和 PyTorch 深度学习框架的可视化炼丹神器!比如看到这行命令,你应该非常熟悉:
代码语言:javascript复制tensorboard --logdir
TAO 工具套件这次更新后,大家可以在 TensorBoard 中可视化很多重要数据(比如训练/验证的损失及变化曲线、模型权重的直方图和预测中间图/结果图等),进而来了解模型的训练性能。
当有了数据可视化,大家就方便更改超参数来对比实验,一步步训练出满意的 AI 模型。
2. 使用 Rest API 将 TAO 工具套件集成到你的应用中
TAO 工具套件现在可以部署为具有 REST API 的服务,这意味着开发人员通过 REST API 构建新的 AI 服务或者将 TAO 工具套件集成到现有服务中,并在 Kubernetes 上管理框架。
3. 通过 BYOM Converter 使用自己的模型权重
BYOM(Bring Your Own Model)是一个基于 Python 的工具包,可以将任何 ONNX 模型转换为与 TAO 工具套件兼容的模型。大家借助 BYOM 工具就可以使用 TAO 来微调和优化自己的模型。比如从 ONNX 导入预训练的权重,进而可以在自己的模型上利用完整的 TAO 功能,例如训练、剪枝和量化等。
注:ONNX 模型的获取方式比较多,你既可以从 ONNX 官网提供的模型库中选择合适的模型,也可以将自己的模型通过 ONNX 工具来转换成 ONNX 格式的模型。
具体来说,TAO BYOM Converter 是提供了一个 CLI 来导入 ONNX 模型并将其转换为 Keras。转换后的模型以基于 EFF 的 .tltb 格式存储。
4. 新增预训练模型
最新版本的 TAO 工具套件中新增了不同领域的代表性模型,涉及 3D 目标检测、动作识别等方向。
PointPillars:3D 目标检测的的代表工作之一,谷歌学术引用量 1100 ,已广泛应用于自动驾驶、机器人等领域。
PoseClassificationNet:将一序列骨架(身体姿态)作为网络输入,并预测这些帧中一个或多个人的动作。最新版本中支持的模型是 ST-GCN 。ST-GCN 可以称为将图卷积神经网络(GCN)应用在基于骨架的动作识别的开山之作 ,谷歌学术引用量1900 ,已广泛应用于视频理解、智能零售等领域。
这次更新的内容还有很多,这里就不一一介绍,感兴趣的同学可以去 TAO 工具套件官网查看。
落地案例介绍
这里再分享一个"智慧农场"的有趣落地案例:OneCup AI 公司利用 NVIDIA TAO 工具套件、DeepStream、Triton 和预训练模型开发了 BETSY(OneCup 的动物人工智能产品),可以帮助牧场主跟踪和监控他们的牲畜(比如奶牛、绵羊等),方便让他们更深入地了解牲畜的健康情况、生长和活动信息。据称在部署的边缘设备上以 2K 分辨率实现高达 30 FPS 推理速度。
总结
如果你想快速上手训练一个模型?
如果你想快速部署应用一个模型?
如果你没有较大的数据集,但想通过迁移学习来获得高度准确的 AI 模型?
那么可以尝试使用 TAO 工具套件!助力你的深度学习科学"炼丹"。
Amusi 期待 TAO 工具套件开放更多的模型,比如大火的 Transformer、Vision Transformer 、NeRF 等,也期待更多新的功能和优化!
附上 TAO 工具套件的资源链接,推荐大家使用!
官网:https://developer.nvidia.com/tao-toolkit
中文官网:https://developer.nvidia.cn/zh-cn/tao-toolkit
TAO 工具套件技术文档:
https://docs.nvidia.com/tao/tao-toolkit/text/overview.html
从 NGC 下载计算机视觉或对话式 AI 模型:
- 计算机视觉模型: https://catalog.ngc.nvidia.com/orgs/nvidia/collections/tltcomputervision
- 对话式 AI 模型: https://catalog.ngc.nvidia.com/orgs/nvidia/collections/tltconversationalai