使用 TAO Deploy 进行分类 (PyTorch) - NVIDIA 文档

为了生成优化的 TensorRT 引擎，分类 (PyTorch) .onnx 文件（首先使用 tao model classification_pyt export 生成）作为 tao deploy classification_pyt gen_trt_engine 的输入。有关训练分类 (PyTorch) 模型的更多信息，请参阅分类 PyTorch 训练文档。在 TAO 5.0.0 中，分类 (PyTorch) 模型不支持 INT8 精度。

将 .onnx 文件转换为 TensorRT 引擎

要转换 .onnx 文件，您可以重用 tao model classification_pyt export 命令中的 spec 文件。

gen_trt_engine

gen_trt_engine 参数定义 TensorRT 引擎生成。

复制
已复制！

            
            gen_trt_engine:
  onnx_file: /path/to/onnx_file
  trt_engine: /path/to/trt_engine
  input_channel: 3
  input_width: 224
  input_height: 224
  tensorrt:
    data_type: fp16
    workspace_size: 1024
    min_batch_size: 1
    opt_batch_size: 16
    max_batch_size: 16

参数	数据类型	默认值	描述	支持的值
`onnx_file`	string		用于 TensorRT 引擎的精度
`trt_engine`	string		TensorRT 引擎的最大工作区大小
`input_channel`	unsigned int	3	输入通道大小。仅支持值 3。	3
`input_width`	unsigned int	224	输入宽度	>0
`input_height`	unsigned int	224	输入高度	>0
`batch_size`	unsigned int	-1	ONNX 模型的批次大小	>=-1
`verbose`	bool	False	启用 TensorRT 日志的详细程度

tensorrt

tensorrt 参数定义 TensorRT 引擎生成。

参数	数据类型	默认值	描述	支持的值
`data_type`	string	fp32	用于 TensorRT 引擎的精度	fp32/fp16/int8
`workspace_size`	unsigned int	1024	TensorRT 引擎的最大工作区大小	>1024
`min_batch_size`	unsigned int	1	用于优化配置文件形状的最小批次大小	>0
`opt_batch_size`	unsigned int	1	用于优化配置文件形状的最佳批次大小	>0
`max_batch_size`	unsigned int	1	用于优化配置文件形状的最大批次大小	>0

使用以下命令运行分类 (PyTorch) 引擎生成

复制
已复制！

            
            tao deploy classification_pyt gen_trt_engine -e /path/to/spec.yaml \
           results_dir=/path/to/results \
           gen_trt_engine.onnx_file=/path/to/onnx/file \
           gen_trt_engine.trt_engine=/path/to/engine/file \
           gen_trt_engine.tensorrt.data_type=<data_type>

必需参数

-e, --experiment_spec：用于设置 TensorRT 引擎生成的实验 spec 文件

可选参数

results_dir：将要转储 JSON 状态日志文件的目录
gen_trt_engine.onnx_file：要转换的 .onnx 模型
gen_trt_engine.trt_engine：将存储生成的引擎的路径
gen_trt_engine.tensorrt.data_type：要导出的精度

示例用法

以下是使用 gen_trt_engine 命令生成 FP16 TensorRT 引擎的示例

复制
已复制！

            
            tao deploy classification_pyt gen_trt_engine -e $DEFAULT_SPEC
           gen_trt_engine.onnx_file=$ONNX_FILE \
           gen_trt_engine.trt_engine=$ENGINE_FILE \
           gen_trt_engine.tensorrt.data_type=FP16

通过 TensorRT 引擎运行评估

您可以重用 TAO 评估 spec 文件，以通过 TensorRT 引擎进行评估。仅当您使用自定义类名称时，才需要 classes 字段。如果未提供此字段，则类映射基于图像文件夹名称的字母数字顺序。以下是示例 spec 文件

复制
已复制！

            
            evaluate:
  trt_engine: /path/to/engine/file
  topk: 1
dataset:
  data:
    samples_per_gpu: 16
    test:
      data_prefix: /raid/ImageNet2012/ImageNet2012/val
      classes: /raid/ImageNet2012/classnames.txt

使用以下命令运行分类 (PyTorch) 引擎评估

复制
已复制！

            
            tao deploy classification_pyt evaluate -e /path/to/spec.yaml \
           results_dir=/path/to/results \
           evaluate.trt_engine=/path/to/engine/file

必需参数

-e, --experiment_spec：用于评估的实验 spec 文件。这应与 tao evaluate spec 文件相同

可选参数

results_dir：将要转储 JSON 状态日志文件和评估结果的目录
evaluate.trt_engine：用于评估的引擎文件

示例用法

以下是使用 evaluate 命令通过 TensorRT 引擎运行评估的示例

复制
已复制！

            
            tao deploy classification_pyt evaluate -e $DEFAULT_SPEC
           results_dir=$RESULTS_DIR \
           evaluate.trt_engine=$ENGINE_FILE

注意

当前，与 TAO PyTorch 评估相比，TAO Deploy trt 评估中，使用 LogisticRegressionHead 的 TAO 分类存在准确性回归。这将在下一个版本中解决。

通过 TensorRT 引擎运行推理

您可以重用 TAO 推理 spec 文件，以通过 TensorRT 引擎进行推理。以下是示例 spec 文件

复制
已复制！

            
            inference:
  trt_engine: /path/to/engine/file
dataset:
  data:
    samples_per_gpu: 16
    test:
      data_prefix: /raid/ImageNet2012/ImageNet2012/val
      classes: /raid/ImageNet2012/classnames.txt

使用以下命令运行分类 (PyTorch) 引擎推理

复制
已复制！

            
            tao deploy classification_pyt inference -e /path/to/spec.yaml \
           results_dir=/path/to/results \
           inference.trt_engine=/path/to/engine/file

必需参数

-e, --experiment_spec：用于推理的实验 spec 文件。这应与 tao inference spec 文件相同。

可选参数

results_dir：将要转储 JSON 状态日志文件和推理结果的目录
inference.trt_engine：用于推理的引擎文件

示例用法

以下是使用 inference 命令通过 TensorRT 引擎运行推理的示例

复制
已复制！

            
            tao deploy classification_pyt inference -e $DEFAULT_SPEC
           results_dir=$RESULTS_DIR \
           evaluate.trt_engine=$ENGINE_FILE

可视化将存储在 $RESULTS_DIR/images_annotated 中，而 KITTI 格式的预测将存储在 $RESULTS_DIR/labels 下。