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
:所需的引擎数据类型。选项为fp32
、fp16
、int8
。默认值为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
下。