使用 Superb AI Suite 和 NVIDIA TAO 工具包创建高质量的计算机视觉应用程序

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

数据标记和模型训练一直被认为是团队在构建 AI / ML 基础设施时面临的最大挑战。这两个步骤都是 ML 应用程序开发过程中的关键步骤，如果做得不当，可能会导致结果不准确和性能下降。要了解更多信息，请参阅人工智能基础设施联盟的2022 年 AI 基础设施生态系统报告。

数据标记对于所有形式的监督学习都至关重要，在监督学习中，整个数据集都被完全标记。它也是半监督学习的一个关键组成部分，半监督学习将一组较小的标记数据与设计用于以编程方式自动标记数据集其余部分的算法相结合。标记对计算机视觉至关重要，计算机视觉是机器学习中最先进和最发达的领域之一。尽管它很重要，但贴标签的速度很慢，因为它需要扩大分布式人力团队的规模。

除了标注之外，模型训练是机器学习的另一个主要瓶颈。训练很慢，因为它需要等待机器完成复杂的计算。它要求团队了解网络、分布式系统、存储、专用处理器（ GPU 或 TPU ）和云管理系统（ Kubernetes 和 Docker ）。

超级 AI 套件与 NVIDIA 〔 TAO 〕工具包

Superb AI 为计算机视觉团队引入了一种新的方法，以大幅减少提供高质量训练数据集所需的时间。团队现在可以通过 Superb AI Suite 来实现这一目标。

图 1 。 Superb AI Suite 为整个数据生命周期提供产品和服务

NVIDIA TAO Toolkit 是建立在 TensorFlow 和 PyTorch 之上的 TAO 框架的低代码版本，它可以抽象掉框架的复杂性，从而加快模型开发过程。TAO 工具包可以让您使用transfer learning，使用自己的数据对 NVIDIA 预训练模型进行微调，并进行推理优化。

图 2 : NVIDIA TAO 工具包 4 . 0 概述

计算机视觉工程师可以将 Superb AI 套件和 TAO 工具包结合使用，以应对数据标记和模型训练的挑战。更具体地说，您可以在 Suite 中快速生成标记数据，并使用 TAO 训练模型，以执行特定的计算机视觉任务，无论是分类、检测还是分割。

准备计算机视觉数据集

本文演示了如何使用 Superb AI Suite 准备一个与 TAO Toolkit 兼容的高质量计算机视觉数据集。它完成了下载数据集、在 Suite 上创建新项目、通过Suite SDK 使用Superb AI 的自动标记功能 快速标记数据集，导出标记的数据集，并设置 TAO Toolkit 配置以使用数据。

步骤 1 ：开始使用 Suite SDK

首先，前往superb-ai.com创建帐户。然后按照快速入门指南安装并认证 Suite CLI。您应该能够安装最新版本的 spb-cli，并检索 Suite 帐户名/访问密钥以进行身份验证。

步骤 2 ：下载数据集

本教程适用于COCO 数据集，一个在计算机视觉研究领域非常受欢迎的大规模目标检测、分割和字幕数据集。

您可以使用此代码片段来下载数据集。将其保存在名为 download-coco.sh 的文件中，然后从终端运行 bash download-coco.sh 。这将创建一个存储 COCO 数据集的 data/ 目录。

下一步是将 COCO 转换为 Suite SDK 格式，以对 COCO 验证 2017 数据集中最常见的五个类进行采样。本教程仅处理边界框注释，但 Suite 也可以处理多边形和关键点。

您可以使用 此代码片段来执行转换。将其保存在名为 convert.py 的文件中，然后从终端运行 Python convert.py 。这将创建一个 upload-info.json 文件，该文件存储有关图像名称和注释的信息。

步骤 3 ：在 Suite SDK 中创建项目

正在使用 Suite SDK 创建项目。本教程使用 Superb AI 创建项目指南。请按照以下配置进行操作。

图 3 。超级人工智能项目创建菜单

1. 选择形象数据类
2. 将项目名称设置为椰子测试
3. 选择注释类型作为边界框
4. 创建五个与 COCO 类名的类名匹配的对象类： "人"、"车"、"椅子"、"书"、"瓶子"

图 4 。在创建过程的这一步，您可以为项目选择和定义对象类

这个过程完成后，您可以查看项目主页面，如图 5 所示。

图 5 。超级 AI 套件主仪表板

步骤 4 ：使用 Suite SDK 上传数据

创建完项目后，开始上传数据。您可以使用此代码片段上传数据。将其保存为upload.py，然后在终端中运行python upload.py --project CocoTest --dataset coco-dataset。

这意味着CocoTest是项目名称，并且coco-dataset是数据集名称。这将启动上传过程，根据设备的处理能力，上传过程可能需要几个小时才能完成。

您可以通过 Suite 网页实时查看上传的数据集，如图 6 所示。

图 6 。通过 Suite 列表视图实时监控上传的数据集

步骤 5 ：标记数据集

下一步是标记 COCO 数据集。要快速完成此操作，请使用 Suite 强大的自动标记功能。更具体地说，自动添加标签 和自定义自动标签都是功能强大的工具，可以通过自动检测对象并对其进行标记来提高标记效率。

自动标签是 Superb AI 开发的一个预训练模型，用于检测和标记 100 多个常见对象，而自定义自动标签是使用您自己的数据训练的模型，用于探测和标记利基对象。

本教程中的 COCO 数据由 Auto Label 能够标记的五种常见对象组成。请按照配置 Auto-Label 的指南进行操作。需要牢记的重要一点是，您需要选择MSCOCO 盒 CAL作为自动标记 AI，并将对象名称与相应应用的对象进行映射。处理 COCO 数据集中的所有 3283 个标签可能需要大约一个小时。

图 7 。创建的自动标签中的对象类设置

自动标签运行完成后，您将看到每个自动标签任务的难度：红色是困难的，黄色是中等的，绿色是容易的。难度越高，"自动标签"错误标记该图像的可能性就越大。

这种难度或估计的不确定性是根据物体尺寸小、照明条件差、场景复杂等因素计算的。在现实世界中，您可以轻松地按难度对标签进行排序和过滤，以便优先处理出错几率较高的标签。

步骤 6 ：从 Suite 导出标记的数据集

获取标记的数据集后，导出并下载标签。标签不仅仅是注释信息，为了充分利用标签来训练 ML 模型，您还需要了解其他信息，例如项目配置和有关原始数据的元信息。要将所有这些信息与注释文件一起下载，请首先请求导出，以便 Suite 系统可以创建一个 zip 文件进行下载。请参阅Suite 中导出和下载标签的指南。

图 8 。通过用户界面导出数据集

导出标签时，将创建一个压缩的 zip 文件供您下载。导出结果文件夹将包含有关整个项目的常规信息、每个标签的注释信息以及每个数据资产的元数据。想要了解更多详情，请访问 导出结果格式 文档。

步骤 7 ：将输出转换为 COCO 格式

接下来，创建一个脚本，将标记的数据转换为可以被 TAO Toolkit 识别的格式，例如COCO 格式。请注意，由于本教程使用了 COCO 数据集，因此数据已经采用了 COCO 格式。例如，您可以在下面找到一个随机导出标签的 JSON 文件：

{
   "objects": [
       {
           "id": "7e9fe8ee-50c7-4d4f-9e2c-145d894a8a26",
           "class_id": "7b8205ef-b251-450c-b628-e6b9cac1a457",
           "class_name": "person",
           "annotation_type": "box",
           "annotation": {
               "multiple": false,
               "coord": {
                   "x": 275.47,
                   "y": 49.27,
                   "width": 86.39999999999998,
                   "height": 102.25
               },
               "meta": {},
               "difficulty": 0,
               "uncertainty": 0.0045
           },
           "properties": []
       },
       {
           "id": "70257635-801f-4cad-856a-ef0fdbfdf613",
           "class_id": "7b8205ef-b251-450c-b628-e6b9cac1a457",
           "class_name": "person",
           "annotation_type": "box",
           "annotation": {
               "multiple": false,
               "coord": {
                   "x": 155.64,
                   "y": 40.61,
                   "width": 98.34,
                   "height": 113.05
               },
               "meta": {},
               "difficulty": 0,
               "uncertainty": 0.0127
           },
           "properties": []
       }
   ],
   "categories": {
       "properties": []
   },
   "difficulty": 1
}

步骤 8 ：准备用于模型训练的标记数据

接下来，使用 SuiteDataset 将 COCO 数据从 Suite 中提取到模型开发中。 SuiteDataset 使 Suite 中导出的数据集可以通过 PyTorch 数据管道访问。下面显示的代码片段实例化了训练集的 SuiteData 对象类。

class SuiteDataset(Dataset):
   """
   Instantiate the SuiteDataset object class for training set
   """

   def __init__(
           self,
           team_name: str,
           access_key: str,
           project_name: str,
           export_name: str,
           train: bool,
           caching_image: bool = True,
           transforms: Optional[List[Callable]] = None,
           category_names: Optional[List[str]] = None,
   ):
       """Function to initialize the object class"""
       super().__init__()

       # Get project setting and export information through the SDK
       # Initialize the Python Client
       client = spb.sdk.Client(team_name=team_name, access_key=access_key, project_name=project_name)
       # Use get_export
       export_info = call_with_retry(client.get_export, name=export_name)
       # Download the export compressed file through download_url in Export
       export_data = call_with_retry(urlopen, export_info.download_url).read()

       # Load the export compressed file into memory
       with ZipFile(BytesIO(export_data), 'r') as export:
           label_files = [f for f in export.namelist() if f.startswith('labels/')]
           label_interface = json.loads(export.open('project.json', 'r').read())
           category_infos = label_interface.get('object_detection', {}).get('object_classes', [])

       cache_dir = None
       if caching_image:
           cache_dir = f'/tmp/{team_name}/{project_name}'
           os.makedirs(cache_dir, exist_ok=True)

       self.client = client
       self.export_data = export_data
       self.categories = [
           {'id': i + 1, 'name': cat['name'], 'type': cat['annotation_type']}
           for i, cat in enumerate(category_infos)
       ]
       self.category_id_map = {cat['id']: i + 1 for i, cat in enumerate(category_infos)}
       self.transforms = build_transforms(train, self.categories, transforms, category_names)
       self.cache_dir = cache_dir

       # Convert label_files to numpy array and use
       self.label_files = np.array(label_files).astype(np.string_)

   def __len__(self):
       """Function to return the number of label files"""
       return len(self.label_files)

   def __getitem__(self, idx):
       """Function to get an item"""
       idx = idx if idx >= 0 else len(self) + idx
       if idx < 0 or idx >= len(self):
           raise IndexError(f'index out of range')

       image_id = idx + 1
       label_file = self.label_files[idx].decode('ascii')

       # Load label information corresponding to idx from the export compressed file into memory
       with ZipFile(BytesIO(self.export_data), 'r') as export:
           label = load_label(export, label_file, self.category_id_map, image_id)

       # Download the image through the Suite sdk based on label_id
       try:
           image = load_image(self.client, label['label_id'], self.cache_dir)
       # Download data in real time using get_data from Suite sdk
       except Exception as e:
           print(f'Failed to load the {idx}-th image due to {repr(e)}, getting {idx + 1}-th data instead')
           return self.__getitem__(idx + 1)

       target = {
           'image_id': image_id,
           'label_id': label['label_id'],
           'annotations': label['annotations'],
       }

       if self.transforms is not None:
           image, target = self.transforms(image, target)
       return image, target

以类似的方式处理测试集。下面的代码片段通过包装 SuiteDataset，使其与 Torchvision COCOEvaluator 配合使用。

class SuiteCocoDataset(C.CocoDetection):
   """
   Instantiate the SuiteCocoDataset object class for test set
   (by wrapping SuiteDataset to make compatible with torchvision's official COCOEvaluator)
   """

   def __init__(
           self,
           team_name: str,
           access_key: str,
           project_name: str,
           export_name: str,
           train: bool,
           caching_image: bool = True,
           transforms: Optional[List[Callable]] = None,
           category_names: Optional[List[str]] = None,
           num_init_workers: int = 20,
   ):
       """Function to initialize the object class"""
       super().__init__(img_folder='', ann_file=None, transforms=None)

       # Call the SuiteDataset class
       dataset = SuiteDataset(
           team_name, access_key, project_name, export_name,
           train=False, transforms=[],
           caching_image=caching_image, category_names=category_names,
       )
       self.client = dataset.client
       self.cache_dir = dataset.cache_dir

       self.coco = build_coco_dataset(dataset, num_init_workers)
       self.ids = list(sorted(self.coco.imgs.keys()))
       self._transforms = build_transforms(train, dataset.categories, transforms, category_names)

   def _load_image(self, id: int):
       """Function to load an image"""
       label_id = self.coco.loadImgs(id)[0]['label_id']
       image = load_image(self.client, label_id, self.cache_dir)
       return image

   def __getitem__(self, idx):
       """Function to get an item"""
       try:
           return super().__getitem__(idx)
       except Exception as e:
           print(f'Failed to load the {idx}-th image due to {repr(e)}, getting {idx + 1}-th data instead')
           return self.__getitem__(idx + 1)

然后可以将 SuiteDataset 和 SuiteCoDataset 用于您的培训代码。下面的代码片段说明了如何使用它们。在模型开发过程中，与train_loader并使用进行评估test_loader.

train_dataset = SuiteDataset(
   team_name=args.team_name,
   access_key=args.access_key,
   project_name=args.project_name,
   export_name=args.train_export_name,
   caching_image=args.caching_image,
   train=True,
)
test_dataset = SuiteCocoDataset(
   team_name=args.team_name,
   access_key=args.access_key,
   project_name=args.project_name,
   export_name=args.test_export_name,
   caching_image=args.caching_image,
   train=False,
   num_init_workers=args.workers,
)

train_loader = DataLoader(
   train_dataset, num_workers=args.workers,
   batch_sampler=G.GroupedBatchSampler(
       RandomSampler(train_dataset),
       G.create_aspect_ratio_groups(train_dataset, k=3),
       args.batch_size,
   ),
   collate_fn=collate_fn,
)
test_loader = DataLoader(
   test_dataset, num_workers=args.workers,
   sampler=SequentialSampler(test_dataset), batch_size=1,
   collate_fn=collate_fn,
)

步骤 9 ：使用 NVIDIA TAO 工具包训练您的模

现在可以使用 Suite 注释的数据来训练对象检测模型。通过使用 TAO 工具包，您可以训练、微调、修剪和导出高度优化和准确的计算机视觉模型，以便将其部署到流行的网络架构和主干网中。在本教程中，您可以选择YOLO v4，它包括在 TAO 中的对象检测模型。

首先，从 TAO 工具包快速入门 下载笔记本示例。

pip3 install nvidia-tao
wget --content-disposition https://api.ngc.nvidia.com/v2/resources/nvidia/tao/tao-getting-started/versions/4.0.1/zip -O getting_started_v4.0.1.zip



$ unzip -u getting_started_v4.0.1.zip  -d ./getting_started_v4.0.1 && rm -rf getting_started_v4.0.1.zip && cd ./getting_started_v4.0.1

接下来，使用以下代码启动笔记本：

$ jupyter notebook --ip 0.0.0.0 --port 8888 --allow-root

在本地主机上打开 Internet 浏览器，然后导航到 URL ：

http://0.0.0.0:8888

要创建 YOLOv4 模型，请打开notebooks/tao_launcher_starter_kit/yolo_v4/yolo_v4.ipynb并按照笔记本上的说明训练模型。

根据结果，不断调整模型，直到达到您的度量目标。如果需要，您可以创建自己的主动学习循环。在真实世界的场景中，查询失败预测的样本，指定人工标注器来注释这批新的样本数据，并用新标注的训练数据补充模型。Superb AI Suite 可以进一步帮助您在后续几轮模型开发中进行数据收集和注释，以不断提高模型性能。

最近发布的 TAO Toolkit 4.0，可以让没有任何人工智能专业知识的用户更容易上手，创建高精度模型。使用 AutoML 自动微调您的超参数，将 TAO Toolkit 交钥匙部署到各种云服务中，将 TAO Toolkit 与第三方 MLOP 服务集成，探索新的基于 transformer 的视觉模型（如 CitySemSegformer 和 Peoplenet Transformer）。

结论

计算机视觉中的数据标注可能会带来许多独特的挑战。由于需要标记的数据量很大，因此该过程可能很困难且成本高昂。此外，该过程可能是主观的，这使得在大型数据集中实现一致的高质量标记输出具有挑战性。

模型训练也可能具有挑战性，因为许多算法和超参数需要调整和优化。这个过程需要对数据和模型有深入的了解，并进行大量实验以获得最佳结果。此外，计算机视觉模型往往需要大量的计算能力来训练，这使得在有限的预算和时间内很难做到这一点。

Superb AI Suite 使您能够收集和标记高质量的计算机视觉数据集。使用 NVIDIA TAO 工具包，您可以优化预先训练的计算机视觉模型。在不牺牲质量的情况下，将两者结合使用可显著加快计算机视觉应用程序的开发时间。

想要更多信息吗？退房时间：

• TAO 工具包 Google Colab 笔记本
• TAO 工具包文档
• Superb AI Suite 标注平台
• Superb AI Suite 文档

关于 Superb AI

Superb AI 提供了一个训练数据平台，使构建、管理和管理计算机视觉数据集比以往任何时候都更快、更容易。我们的解决方案专注于标签和质量保证的适应性自动化模型，帮助公司大幅减少为计算机视觉模型构建数据管道所需的时间和成本。2018 年，由拥有数十年经验的研究人员和工程师（包括 25 多项出版物、7300 多项引用和 100 多项专利）发起，我们的愿景是让处于各个阶段的公司能够比以往更快地开发计算机视觉应用程序。

Superb AI 也是通过 NVIDIA Inception Program for Startups 由 NVIDIA 提供的计划，有助于培育世界尖端初创公司的发展，为他们提供 NVIDIA 技术和专家，与风险投资家建立联系的机会，以及提高知名度的市场营销支持。

阅读原文