TAO v5.5.0

UNet 与 TAO Deploy

tao model export 生成的 UNet .onnx 文件被用作 tao deploy 的输入,以生成优化的 TensorRT 引擎。有关训练 UNet 的更多信息,请参阅 UNet 训练文档

可以使用与 tao model unet export 命令相同的 spec 文件。

使用以下命令运行 UNet 引擎生成

复制
已复制!
            

tao deploy unet gen_trt_engine [-h] [-v] -m MODEL_PATH -e EXPERIMENT_SPEC -r RESULTS_DIR [-k KEY] [--data_type {fp32,fp16,int8}] [--engine_file ENGINE_FILE] [--cal_image_dir CAL_IMAGE_DIR] [--cal_data_file CAL_DATA_FILE] [--cal_cache_file CAL_CACHE_FILE] [--cal_json_file CAL_JSON_FILE] [--max_batch_size MAX_BATCH_SIZE] [--min_batch_size MIN_BATCH_SIZE] [--opt_batch_size OPT_BATCH_SIZE] [--batch_size BATCH_SIZE] [--batches BATCHES] [--max_workspace_size MAX_WORKSPACE_SIZE] [-s STRICT_TYPE_CONSTRAINTS] [--force_ptq FORCE_PTQ] [--gpu_index GPU_INDEX] [--log_file LOG_FILE]

必需参数

  • -m, --model_path:要转换的 .onnx.etlt 模型。

  • -e, --experiment_spec:用于设置 TensorRT 引擎生成的实验 spec 文件。这应与导出规范文件相同。

  • -r, --results_dir:将要转储 JSON 状态日志文件的目录

可选参数

  • -h, --help:显示此帮助消息并退出。

  • -k, --key:用于加载 .etlt 模型的用户特定编码密钥。

  • --data_type:所需的引擎数据类型。选项为 fp32fp16int8。默认值为 fp32。校准缓存将在 INT8 模式下生成。如果使用 INT8,则需要以下 INT8 参数。

  • --engine_file:序列化的 TensorRT 引擎文件的路径。请注意,此文件是硬件特定的,不能在 GPU 之间通用。由于 TensorRT 引擎文件是硬件特定的,因此除非部署 GPU 与训练 GPU 完全相同,否则您不能使用此引擎文件进行部署。

  • -s, --strict_type_constraints:一个布尔标志,指示在构建 TensorRT 引擎时是否应用 TensorRT 严格类型约束。

  • --gpu_index:用于导出模型的(离散)GPU 的索引。如果机器安装了多个 GPU,您可以指定要运行导出的 GPU 索引。请注意,gen_trt_engine 只能在单个 GPU 上运行。

  • --log_file:日志文件的路径。默认路径为“stdout”。

INT8 引擎生成必需参数

  • --cal_data_file:为校准引擎生成的 Tensorfile。如果与 --cal_image_dir 一起使用,它也可以是输出文件。

注意

--cal_image_dir 参数用于图像,并应用必要的预处理以在 --cal_data_file 参数中提到的路径生成 tensorfile,而 tensorfile 又用于校准。tensorfile 中批次的数量从设置为 --batches 参数的值中获得,而 batch_size 从设置为 --batch_size 参数的值中获得。请确保 --cal_image_dir 中提到的目录至少包含 batch_size * batches 张图像。有效图像扩展名为 .jpg、.jpeg 和 .png。在这种情况下,校准张量的 input_dimensions.etlt 模型的输入层派生而来。

INT8 引擎生成可选参数

  • --cal_image_dir:用于校准的图像目录。如果未指定,将使用 spec 文件中的训练数据。

  • --cal_cache_file:保存校准缓存文件的路径。默认值为 ./cal.bin

  • --cal_json_file:包含 QAT 模型张量比例的 json 文件的路径。如果要生成 QAT 模型的引擎,则此参数是必需的。

  • --batches:用于校准的批次数量。默认值为 10。

  • --batch_size:用于校准的批次大小。默认值为 1。

  • --max_batch_size:TensorRT 引擎的最大批次大小。默认值为 1。

  • --min_batch_size:TensorRT 引擎的最小批次大小。默认值为 1。

  • --opt_batch_size:TensorRT 引擎的最佳批次大小。默认值为 1。

  • --max_workspace_size:TensorRT 引擎的最大工作区大小,单位为 Gb。默认值为:(2 Gb)。

  • --force_ptq:一个布尔标志,用于强制对导出的 etlt 模型进行后训练量化。

注意

当为启用 QAT 训练的模型生成 TensorRT 引擎时,需要由 cal_cache_file 参数定义的张量比例因子。但是,请注意,当前版本的 QAT 不原生支持 Jetson 中的 DLA int8 部署。为了在带有 DLA int8 的 Jetson 上部署此模型,请使用 --force_ptq 标志以使用 TensorRT 后训练量化来生成校准缓存文件。

示例用法

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

复制
已复制!
            

tao deploy unet gen_trt_engine -m /workspace/unet.onnx \ -e /workspace/default_spec.txt \ -r /export/ \ --cal_image_dir /workspace/tao-experiments/data/isbi/images/val \ --data_type int8 \ --batch_size 8 \ --batches 10 \ --cal_cache_file /export/cal.bin \ --cal_data_file /export/cal.tensorfile \ --engine_file /export/int8.engine

与 TAO 评估 spec 文件相同的 spec 文件。示例 spec 文件

复制
已复制!
            

model_config { model_input_width: 320 model_input_height: 320 model_input_channels: 1 num_layers: 18 all_projections: true load_graph: true arch: "resnet" use_batch_norm: False training_precision { backend_floatx: FLOAT32 } } dataset_config { dataset: "custom" augment: False input_image_type: "grayscale" train_images_path:"/workspace/tao-experiments/data/isbi/images/train" train_masks_path:"/workspace/tao-experiments/data/isbi/masks/train" val_images_path:"/workspace/tao-experiments/data/isbi/images/val" val_masks_path:"/workspace/tao-experiments/data/isbi/masks/val" test_images_path:"/workspace/tao-experiments/data/isbi/images/test" data_class_config { target_classes { name: "foreground" mapping_class: "foreground" label_id: 0 } target_classes { name: "background" mapping_class: "background" label_id: 1 } } }

使用以下命令运行 UNet 引擎评估

复制
已复制!
            

tao deploy unet evaluate [-h] -e EXPERIMENT_SPEC -m MODEL_PATH -r RESULTS_DIR [-i IMAGE_DIR] [-l LABEL_DIR] [-b BATCH_SIZE] [--gpu_index GPU_INDEX] [--log_file LOG_FILE]

必需参数

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

  • -m, --model_path:要运行评估的引擎文件。

  • -r, --results_dir:将存储评估结果的目录

可选参数

  • -i, --image_dir:测试图像所在的目录。如果未指定,将使用 spec 文件中的 dataset_config.val_data_sources.data_source.images_path

  • -l, --label_dir:测试标注所在的目录。如果未指定,将使用 spec 文件中的 dataset_config.val_data_sources.data_source.labels_path

  • -b, --batch_size:用于评估的批次大小。请注意,此值不能大于引擎生成期间使用的 --max_batch_size。如果未指定,将使用 training_config.batch_size 代替。

示例用法

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

复制
已复制!
            

tao deploy unet evaluate -m /export/int8.engine \ -e /workspace/default_spec.txt \ -i /workspace/tao-experiments/data/isbi/images/val \ -l /workspace/tao-experiments/data/isbi/masks/val \ -r /workspace/tao-experiments/evaluate

复制
已复制!
            

tao deploy unet inference [-h] -e EXPERIMENT_SPEC -m MODEL_PATH -r RESULTS_DIR [-i IMAGE_DIR] [-b BATCH_SIZE] [--gpu_index GPU_INDEX] [--log_file LOG_FILE]

必需参数

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

  • -m, --model_path:要运行评估的引擎文件。

  • -r, --results_dir:将存储评估结果的目录

可选参数

  • -i, --image_dir:测试图像所在的目录。如果未指定,将使用 spec 文件中的 dataset_config.val_data_sources.data_source.images_path

  • -b, --batch_size:用于评估的批次大小。请注意,此值不能大于引擎生成期间使用的 --max_batch_size。如果未指定,将使用 training_config.batch_size 代替。

示例用法

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

复制
已复制!
            

tao deploy unet inference -m /export/int8.engine \ -e /workspace/default_spec.txt \ -i /workspace/tao-experiments/data/isbi/images/val \ -r /workspace/tao-experiments/inference

可视化结果将存储在 $RESULTS_DIR/images_annotated 下,PNG 格式的掩码预测将存储在 $RESULTS_DIR/labels 下。

上一个 Segformer 与 TAO Deploy
下一个 YOLOv3 与 TAO Deploy
© 版权所有 2024, NVIDIA。 上次更新于 2024 年 10 月 15 日。