优化训练流程

TAO 支持的所有深度神经网络任务都提供了一个 train 命令，使用户能够训练模型。训练可以在一个或多个 GPU 上完成。NVIDIA TAO 提供了一个简单的命令行界面，用于训练用于分类、对象检测和实例分割的深度学习模型。为了加快训练过程，train 命令支持多 GPU 训练。您可以使用 --gpus N 选项调用多 GPU 训练会话，其中 N 是您要使用的 GPU 数量。N 必须小于给定节点中可用于训练的 GPU 数量。

以下优化也包含在 train 命令中

• 量化感知训练 (QAT)

• 自动混合精度 (AMP)

TAO 为其对象检测网络（即 DetectNet_v2、SSD、DSSD、YOLOv3、YOLOv4、YOLOv4-tiny、RetinaNet、FasterRCNN）和语义分割网络（即 UNet）支持量化感知训练 (QAT)。量化感知训练在训练模型时模拟推理时量化，该模型随后可能被下游推理平台用于生成实际的量化模型。在训练期间，对权重和张量量化为 INT8 产生的误差进行建模，从而使模型能够适应并减轻误差。在 QAT 期间，训练图中构建的模型被修改为

1. 用支持权重伪量化的节点替换现有节点。

2. 将现有激活转换为 ReLU-6（输出节点除外）。

3. 添加量化和反量化 (QDQ) 节点以计算中间张量的动态范围。

训练期间计算的动态范围被序列化到 export 的缓存文件中，然后 TensorRT 可以解析该文件以创建优化的推理引擎。要在训练期间启用 QAT，只需在每个受支持网络的相应 spec 文件的 training_config 字段中将 enable_qat 参数设置为 true。与传统的基于校准的 INT8 TensorRT 推理相比，QAT 训练的优势通常是在使用 TensorRT 进行 INT8 推理时具有更好的准确性。

下图描述了训练量化感知模型的推荐工作流程。

TAO 现在支持自动混合精度 (AMP) 训练。DNN 训练传统上依赖于使用 IEEE 单精度格式训练其张量。然而，通过混合精度训练，可以在训练图中混合使用 FP16 和 FP32 操作，以帮助加速训练，同时又不影响准确性。使用 AMP 有以下几个好处

• 加速数学密集型操作，例如线性和卷积层。

• 通过访问单精度一半的字节数来加速内存限制的操作

• 减少训练模型的内存需求，从而实现更大的模型或更大的小批量。

在 TAO 中，启用 AMP 非常简单，只需在运行 train 命令时在命令行中设置 --use_amp 标志即可。这将通过使用 FP16 张量核心来帮助加速训练。请注意，AMP 仅在具有 Volta 或更高架构的 GPU 上受支持。

TAO 现在支持 DINO 的知识蒸馏。知识蒸馏是一种技术，其中训练小型学生模型来模仿大型教师模型的行为。教师模型用于为学生模型提供软目标，这有助于学生模型更好地学习。训练学生模型以最大程度地减少教师模型和学生模型的预测之间的差异。当教师模型很大且运行计算成本很高，而学生模型很小且可以部署在边缘设备上时，此技术很有用。

在 DINO 训练中，我们可以通过在配置规范中设置 distill 字段来执行知识蒸馏，其中包含有关教师模型和蒸馏损失绑定的详细信息，并使用 distill 入口点而不是 train 入口点。

            

            
distill:
  teacher:
     backbone: fan_small
     train_backbone: False
     num_feature_levels: 4
     dec_layers: 6
     enc_layers: 6
     num_queries: 900
     dropout_ratio: 0.0
     dim_feedforward: 2048
  pretrained_teacher_model_path: /workspace/tao-experiments/dino/pretrained_dino_coco_vdino_fan_small_trainable_v1.0/dino_fan_small_ep12.pth
  bindings:
  - teacher_module_name: 'model.backbone.0.body'
     student_module_name: 'model.backbone.0.body'
     criterion: L2
     weight: 1.0
        

            

            
tao model dino distill -e /path/to/spec.yaml