OCDNet - NVIDIA 文档 - NVIDIA 文档

OCDNet 是 TAO 中包含的光学字符检测模型。它支持以下任务

训练
评估
推理
剪枝
导出

这些任务可以使用 TAO 启动器从命令行调用，使用以下约定

复制
已复制!

            
            tao model ocdnet <sub_task> <args_per_subtask>

其中 args_per_subtask 是给定子任务所需的命令行参数。以下各节详细解释了每个子任务。

准备数据集

OCDNet 的数据集包含图像和相应的标签文件。

训练数据集和测试数据集都必须遵循相同的结构。目录结构应组织如下，其中图像的目录名称为 img，标签文件的目录名称为 gt。默认情况下，标签文件应使用 gt_ 作为前缀，以便与相应的图像文件进行比较。

确切的目录名称 train 和 test 不是必需的，但按惯例是首选的。

复制
已复制!

            
            /train
  /img
    img_0.jpg
    img_1.jpg
    ...
  /gt
    gt_img_0.txt
    gt_img_1.txt
    ...
/test
  /img
    img_0.jpg
    img_1.jpg
    ...
  /gt
    gt_img_0.txt
    gt_img_1.txt
    ...

下面是来自公共 ICDAR2015 数据集的标签文件示例

复制
已复制!

            
            $ cat ICDAR2015/test/gt/gt_img_14.txt
  268,82,335,93,332,164,267,164,the
  344,94,433,112,427,159,336,163,Future
  208,191,374,184,371,213,208,241,Communications
  370,176,420,176,416,204,373,213,###
  1,57,261,76,261,187,0,190,venting
  1,208,203,200,203,241,3,294,ntelligence.

注意

标签文件包含所有点的坐标。最后一个是文本。如果文本是 ### 并且训练规范文件将 ignore_tags 设置为 ['###']，则在训练期间将忽略这些行。

创建实验规范文件

OCDNet 的规范文件包括 model、train、dataset 和 evaluate，以及其他全局参数。以下是使用 FAN-tiny 主干在 ICDAR2015 数据集上训练 OCDNet 模型的规范文件示例。

复制
已复制!

            
            model:
  load_pruned_graph: False
  pruned_graph_path: '/results/prune/pruned_0.1.pth'
  pretrained_model_path: '/data/ocdnet/ocdnet_fan_tiny_2x_icdar.pth'
  backbone: fan_tiny_8_p4_hybrid
  enlarge_feature_map_size: True
  activation_checkpoint: True
train:
  num_gpus: 1
  results_dir: /results/train
  num_epochs: 10
  resume_training_checkpoint_path: '/results/train/resume.pth'
  checkpoint_interval: 5
  validation_interval: 5
  seed: 1234
  is_dry_run: False
  precision: fp32
  model_ema: False
  model_ema_decay: 0.999
  trainer:
    clip_grad_norm: 5.0
  optimizer:
    type: Adam
    args:
      lr: 0.001
  lr_scheduler:
    type: WarmupPolyLR
    args:
      warmup_epoch: 3
  post_processing:
    type: SegDetectorRepresenter
    args:
      thresh: 0.3
      box_thresh: 0.55
      max_candidates: 1000
      unclip_ratio: 1.5
  metric:
    type: QuadMetric
    args:
      is_output_polygon: false
dataset:
  train_dataset:
    data_name: ICDAR2015Dataset
    data_path: ['/data/ocdnet_vit/train']
    args:
      pre_processes:
        - type: IaaAugment
          args:
            - {'type':Fliplr, 'args':{'p':0.5}}
            - {'type': Affine, 'args':{'rotate':[-45,45]}}
            - {'type':Sometimes,'args':{'p':0.2, 'then_list':{'type': GaussianBlur, 'args':{'sigma':[1.5,2.5]}}}}
            - {'type':Resize,'args':{'size':[0.5,3]}}
        - type: EastRandomCropData
          args:
            size: [640,640]
            max_tries: 50
            keep_ratio: true
        - type: MakeBorderMap
          args:
            shrink_ratio: 0.4
            thresh_min: 0.3
            thresh_max: 0.7
        - type: MakeShrinkMap
          args:
            shrink_ratio: 0.4
            min_text_size: 8

      img_mode: BGR
      filter_keys: [img_path,img_name,text_polys,texts,ignore_tags,shape]
      ignore_tags: ['*', '###']
    loader:
      batch_size: 1
      pin_memory: true
      num_workers: 12
  validate_dataset:
    data_name: ICDAR2015Dataset
    data_path: ['/data/ocdnet_vit/test']
    args:
      pre_processes:
        - type: Resize2D
          args:
            short_size:
              - 1280
              - 736
            resize_text_polys: true
      img_mode: BGR
      filter_keys: []
      ignore_tags: ['*', '###']
    loader:
      batch_size: 1
      pin_memory: false
      num_workers: 1

下表提供了规范文件的顶层描述。

参数	数据类型	默认值	描述	支持的值
`model`	dict 配置	–	模型架构的配置
`dataset`	dict 配置	–	数据集的配置
`训练`	dict 配置	–	训练任务的配置
`评估`	dict 配置	–	评估任务的配置
`推理`	dict 配置	–	推理任务的配置
`encryption_key`	string	无	用于加密和解密模型文件的加密密钥
`results_dir`	string	/results	保存实验结果的目录
`导出`	dict 配置	–	ONNX 导出任务的配置
`gen_trt_engine`	dict 配置	–	TensorRT 生成任务的配置。仅在 TAO 部署中使用
`剪枝`	dict 配置	–	剪枝任务的配置
`name`	str	–

模型

model 参数提供模型的参数列表。

参数	数据类型	默认值	描述	支持的值
`load_pruned_graph`	bool	`false`	一个标志，指定是否加载剪枝图。如果正在对剪枝模型执行训练/评估/导出/推理，则设置为 True。	true/false
`pruned_graph_path`	string	–	剪枝图模型的路径（如果 `load_pruned_graph` 为 True）	unix 路径
`pretrained_model_path`	string	–	预训练模型的路径	unix 路径
`backbone`	string	deformable_resnet18	模型的主干	deformable_resnet18 deformable_resnet50 fan_tiny_8_p4_hybrid
`enlarge_feature_map_size`	bool	`false`	一个标志，指定是否扩大 FAN-tiny 主干的输出特征图大小。当使用 `deformable_resnet` 主干时，此标志无效。	true/false
`activation_checkpoint`	bool	`false`	一个标志，指定是否使用激活检查点来节省 GPU 内存。当使用 `deformable_resnet` 主干时，此标志无效。	true/false true/false

训练

train 参数提供训练的参数。

参数	数据类型	默认值	描述	支持的值
`num_gpus`	unsigned int	1	用于分布式训练的 GPU 数量	>0
`gpu_ids`	List[int]	[0]	用于分布式训练的 GPU 索引
`seed`	unsigned int	1234	用于 random、NumPy 和 torch 的随机种子	>0
`num_epochs`	unsigned int	10	运行实验的总 epoch 数	>0
`checkpoint_interval`	unsigned int	1	保存检查点的 epoch 间隔	>0
`validation_interval`	unsigned int	1	运行验证的 epoch 间隔	>0
`resume_training_checkpoint_path`	string		要从中恢复训练的中间 PyTorch Lightning 检查点
`results_dir`	string	/results/train	保存训练结果的目录
`optimizer`	dict 配置	–	优化器的配置	–
`lr_scheduler`	dict 配置	–	lr_scheduler 的配置	–
`post_processing`	dict 配置	–	post_processing 的配置。	–
`metric`	dict 配置	–	用于度量计算的配置。支持 QuadMetric。如果 `is_output_polygon` 为 True，将生成多边形。如果为 False，将生成 BBox。	–
`is_dry_run`	bool	`false`	如果此标志为 True，则仅运行一个批次。仅建议将此标志用于调试目的。	true/false
`precision`	string	fp32	模型将训练的精度。如果此值设置为“fp16”，则将启用 AMP 训练	fp32/fp16
`model_ema`	bool	`false`	一个标志，用于启用模型 EMA。默认值为 False。如果值为 True，则在训练期间将启用模型 EMA	true/false
`model_ema_decay`	float	0.999	模型 EMA 的衰减。默认值为 0.999。仅当 `model_ema` 设置为 True 时才使用此值。	(0, 1]

optimizer

复制
已复制!

            
            optimizer:
  type: Adam
  args:
    lr: 0.001

参数	数据类型	默认值	描述	支持的值
`type`	string	Adam	优化器类型	Adam
`lr`	float	–	初始学习率	>=0.0

lr_scheduler

复制
已复制!

            
            lr_scheduler:
  type: WarmupPolyLR
  args:
    warmup_epoch: 3

参数	数据类型	默认值	描述	支持的值
`type`	string	WarmupPolyLR	通过多项式函数衰减学习率。学习率在预热阶段增加到初始值，并在训练阶段从初始值减小到零。	WarmupPolyLR
`warmup_epoch`	unsigned int	3	预热 epoch，学习率在此期间增加到初始值（即 `optimizer.args.lr`）。预热 epoch 不应与 `num_epochs` 相同。	>=0

post_processing

复制
已复制!

            
            post_processing:
  type: SegDetectorRepresenter
  args:
    thresh: 0.3
    box_thresh: 0.55
    max_candidates: 1000
    unclip_ratio: 1.5

参数	数据类型	默认值	描述	支持的值
`type`	string	SegDetectorRepresenter	post_processing 的名称。post_processing 将生成 BBox 或多边形。	SegDetectorRepresenter
`thresh`	float	0.3	二值化的阈值，用于生成近似二值图。	0.0 ~ 1.0
`box_thresh`	float	0.7	BBox 阈值。如果有效面积低于此阈值，则将忽略预测，这意味着未检测到文本。	0.0 ~ 1.0
`max_candidates`	unsigned int	1000	最大候选输出。如果在一个区域检测到字符，但在图像的另一个区域未检测到字符，请增大此参数。	> 1
`unclip_ratio`	float	1.5	在概率图中，使用 Vatti 裁剪算法的 unclip 比率。如果此比率设置得较大，则 BBox 看起来会更大。	>0.0

数据集

数据集由两个部分定义：train_dataset 和 validate_dataset

参数	数据类型	默认值	描述	支持的值
`train_dataset`	dict 配置	–	训练数据集的配置	–
`validate_dataset`	dict 配置	–	验证数据集的配置	–

下面提供了 train_dataset 的参数。

参数	数据类型	默认值	描述	支持的值
`data_name`	string	ICDAR2015Dataset	数据集名称。对于“ICDAR2015Dataset”，标签文件应使用 `gt_` 作为前缀。对于“UberDataset”，标签文件应使用 `truth_` 作为前缀。	ICDAR2015Dataset UberDataset
`data_path`	string list	–	包含用于训练的图像的路径列表：例如，`['path_1']` 或 `['path_1', 'path_2', ...]`	–
`pre_processes`	dict	–	预处理配置（有关更多详细信息，请参阅）train_preprocess	–
`img_mode`	string	BGR	图像模式	BGR、RGB、GRAY
`filter_keys`	string list	`['img_path', 'img_name', 'text_polys', 'texts', 'ignore_tags', 'shape']`	要忽略的键	–
`ignore_tags`	string list	`['*', '###']`	不用于训练的标签	–
`batch_size`	unsigned int	False	批次大小。如果遇到内存不足错误，请设置为较低的值。	>0
`pin_memory`	bool	False	一个标志，指定是否启用固定内存	true/false
`num_workers`	unsigned int	1	用于加载数据的线程数	>=0

train_preprocess

复制
已复制!

            
            pre_processes:
  - type: IaaAugment
    args:
      - {'type':Fliplr, 'args':{'p':0.5}}
      - {'type': Affine, 'args':{'rotate':[-45,45]}}
      - {'type':Sometimes,'args':{'p':0.2, 'then_list':{'type': GaussianBlur, 'args':{'sigma':[1.5,2.5]}}}}
      - {'type':Resize,'args':{'size':[0.5,3]}}
  - type: EastRandomCropData
    args:
      size: [640,640]
      max_tries: 50
      keep_ratio: true
  - type: MakeBorderMap
    args:
      shrink_ratio: 0.4
      thresh_min: 0.3
      thresh_max: 0.7
  - type: MakeShrinkMap
    args:
      shrink_ratio: 0.4
      min_text_size: 8

参数	数据类型	默认值	描述	支持的值
`IaaAugment`	dict list	`{'type':Fliplr, 'args':{'p':0.5}}` `{'type': Affine, 'args':{'rotate':[-10,10]}}` `{'type':Sometimes,'args':{'p':1.0, 'then_list':{'type': GaussianBlur, 'args':{'sigma':[1.5,2.5]}}}}` `{'type':Resize,'args':{'size':[0.5,3]}}`	使用 imgaug 执行增强。“Fliplr”、“Affine”、“Sometimes”、“GaussianBlur”和“Resize”默认使用。 `p` 定义每个图像被翻转的概率。`rotate` 定义按随机值旋转图像时的度数范围。 `Sometimes` 仅定义具有一个或多个增强器的所有图像的 `p` 百分比。`then_list` 定义要应用于 `p` 百分比的所有图像的增强器 `GaussianBlur` 使用高斯核定义模糊。`sigma` 定义高斯核的标准偏差。 `size` 定义调整每个图像大小（与其原始大小相比）时的范围。	`p`: 0.0 ~ 1.0 `roate`: -180 ~ 180 `sigma`: -180 ~ 180 `resize`: >0.0 ~ >0.0
`EastRandomCropData`	dict 配置	– –	增强后的随机裁剪。`size` 定义裁剪的目标大小（宽度、高度）。宽度和高度应为 32 的倍数。`max_tries` 定义尝试裁剪的最大次数，因为裁剪区域可能太小或裁剪可能失败。`keep_ratio` 指定是否保持宽高比。	`size`: [>0, >0] `max_tries`: >0 `keep_ratio`: true/false
`MakeBorderMap`	dict 配置	–	定义生成阈值图时的参数。`shrink_ratio` 用于计算扩展/收缩多边形与原始文本多边形之间的距离。`thresh_min` 和 `thresh_max` 将设置生成阈值图时的阈值范围。	0.0 ~ 1.0
`MakeShrinkMap`	dict 配置	–	定义生成概率图时的参数。`shrink_ratio` 用于生成收缩的多边形。`min_text_size` 指定如果文本的高度或宽度低于此参数，则将忽略该文本。	0.0 ~ 1.0

validate_dataset 的参数与 train_dataset 类似，除了下面的 validation_preprocess。

validation_preprocess

复制
已复制!

            
            pre_processes:
  - type: Resize2D
    args:
      short_size:
        - 1280
        - 736
      resize_text_polys: true

参数	数据类型	默认值	描述	支持的值
`type`	string	Resize2D	在评估之前调整图像和标签的大小。	Resize2D
`short_size`	list	–	将图像调整为（宽度 x 高度）。	>0、>0 和 32 的倍数。
`resize_text_polys`	bool	–	一个标志，指定是否调整文本坐标的大小	true/false

评估

以下是在 ICDAR2015 数据集上进行评估的规范文件示例

复制
已复制!

            
            model:
  load_pruned_graph: False
  pruned_graph_path: '/results/prune/pruned_0.1.pth'
  backbone: deformable_resnet18
evaluate:
  results_dir: /results/evaluate
  checkpoint: /results/train/model_best.pth
  gpu_id: 0
  post_processing:
    type: SegDetectorRepresenter
    args:
      box_thresh: 0.55
      max_candidates: 1000
      unclip_ratio: 1.5
  metric:
    type: QuadMetric
    args:
      is_output_polygon: false
dataset:
  validate_dataset:
    data_path: ['/data/ocdnet/test']
    args:
      pre_processes:
        - type: Resize2D
          args:
            short_size:
              - 1280
              - 736
            resize_text_polys: true
      img_mode: BGR
      filter_keys: []
      ignore_tags: ['*', '###']
    loader:
      batch_size: 1
      shuffle: false
      pin_memory: false
      num_workers: 4

推理

以下是运行推理的规范文件示例

复制
已复制!

            
            model:
  load_pruned_graph: false
  pruned_graph_path: '/results/prune/pruned_0.1.pth'
  backbone: deformable_resnet18
inference:
  checkpoint: '/results/train/model_best.pth'
  input_folder: /data/ocdnet/test/img
  width: 1280
  height: 736
  img_mode: BGR
  polygon: false
  results_dir: /results/inference
  post_processing:
    type: SegDetectorRepresenter
    args:
      thresh: 0.3
      box_thresh: 0.55
      max_candidates: 1000
      unclip_ratio: 1.5

inference 参数定义推理过程的超参数。推理绘制边界框或多边形，并在图像中可视化。

参数	数据类型	默认值	描述	支持的值
`checkpoint`	string	–	pth 模型的路径	Unix 路径
`results_dir`	string	/results/inference	保存推理结果的目录
`num_gpus`	unsigned int	1	用于分布式推理的 GPU 数量	>0
`gpu_ids`	List[int]	[0]	用于分布式推理的 GPU 索引
`input_folder`	string	–	推理的输入文件夹的路径	Unix 路径
`width`	unsigned int	–	输入宽度	>=1
`height`	unsigned int	–	输入高度	>=1
`img_mode`	string	–	图像模式	BGR/RGB/GRAY
`polygon`	bool	–	True 值指定 BBox，False 值指定多边形。	true, false

训练模型

使用以下命令运行 OCDnet 训练

复制
已复制!

            
            tao model ocdnet train -e <experiment_spec_file>
                 [results_dir=<global_results_dir>]
                 [model.<model_option>=<model_option_value>]
                 [dataset.<dataset_option>=<dataset_option_value>]
                 [train.<train_option>=<train_option_value>]
                 [train.gpu_ids=<gpu indices>]
                 [train.num_gpus=<number of gpus>]

必需参数

唯一必需的参数是实验规范的路径

-e, --experiment_spec：用于设置训练实验的实验规范文件

可选参数

您可以设置可选参数来覆盖实验规范文件中的选项值。

-h, --help：显示此帮助消息并退出。
model.<model_option>：模型选项。
dataset.<dataset_option>：数据集选项。
train.<train_option>：训练选项。
train.optim.<optim_option>：优化器选项

注意

对于训练、评估和推理，我们为每个相应的任务公开 2 个变量：num_gpus 和 gpu_ids，它们分别默认为 1 和 [0]。如果两者都已传递，但彼此不一致，例如 num_gpus = 1，gpu_ids = [0, 1]，则它们将被修改为遵循具有更多 GPU 的设置，例如 num_gpus = 1 -> num_gpus = 2。

检查点和恢复训练

在每个 train.checkpoint_interval，都会保存 PyTorch Lightning 检查点。它被称为 model_epoch_<epoch_num>.pth。这些保存在 train.results_dir 中，如下所示

复制
已复制!

            
            $ ls /results/train

'model_epoch_000.pth'
'model_epoch_001.pth'
'model_epoch_002.pth'
'model_epoch_003.pth'
'model_epoch_004.pth'

最新的检查点也将另存为 ocd_model_latest.pth。如果 ocd_model_latest.pth 存在于 train.results_dir 中，则训练会自动从 ocd_model_latest.pth 恢复。如果提供了 train.resume_training_checkpoint_path，则它将被 train.resume_training_checkpoint_path 取代。

此逻辑的主要含义是，如果您希望从头开始触发全新训练，则可以

注意

默认情况下，训练使用 DDP（分布式数据并行）策略。当使用多 GPU 进行训练时，仅当评估图像是 num_gpus * evaluate_batch_size 的倍数时，训练期间的 hmean 结果才会与运行 tao model ocdnet evaluate 的 hmean 结果相同。

评估模型

使用以下命令运行 OCDNet 评估

复制
已复制!

            
            tao model ocdnet evaluate  -e <experiment_spec_file>
                 evaluate.checkpoint=<model to be evaluated>
                 [evaluate.<evaluate_option>=<evaluate_option_value>]
                 [evaluate.gpu_ids=<gpu indices>]
                 [evaluate.num_gpus=<number of gpus>]

必需参数

-e, --experiment_spec：用于设置评估实验的实验规范文件。
evaluate.checkpoint：要评估的 .pth 模型。

可选参数

evaluate.<evaluate_option>：评估选项。

在 OCDNet 模型上运行推理

复制
已复制!

            
            tao ocdnet inference -e <experiment_spec_file>
           inference.checkpoint=<model to be inferenced>
           inference.input_folder=<path to input folder>
           [inference.<inference_option>=<inference_option_value>]
           [inference.gpu_ids=<gpu indices>]
           [inference.num_gpus=<number of gpus>]

必需参数

-e, --experiment_spec：用于设置推理实验的实验规范文件。
inference.checkpoint：用于推理的 .pth 模型。
inference.input_folder：输入文件夹的路径。

可选参数

inference.<inference_option>：推理选项。

注意

推理需要 gt 文件夹中已存在的标签文件。如果没有标签文件，请在 gt 文件夹下生成虚拟标签。使用以下脚本作为参考

复制
已复制!

            
            #!/bin/bash
folder_path=/workspace/datasets/ICDAR2015/datasets/test
mkdir -p ${folder_path}/gt
for filename in `ls ${folder_path}/img`; do
    touch "${folder_path}/gt/gt_${filename%.*}.txt"
    echo "10,10,10,20,20,10,20,20,###" > "${folder_path}/gt/gt_${filename%.*}.txt"
done

剪枝和重新训练 OCDNet 模型

模型剪枝减少模型参数以提高推理帧率 (FPS)，同时保持几乎相同的 hmean。

剪枝应用于已训练的 OCDNet 模型。剪枝后，将生成剪枝图模型。它是一个参数较少的新模型。获得此剪枝图模型后，您必须在同一数据集上重新训练它以恢复 hmean。在重新训练期间，您需要启用加载此剪枝图模型并设置此模型的路径。

prune 参数定义剪枝过程的超参数。

复制
已复制!

            
            prune:
  checkpoint: /results/train/model_best.pth
  ch_sparsity: 0.2
  round_to: 32
  p: 2
  results_dir: /results/prune
  verbose: True

model:
  backbone: fan_tiny_8_p4_hybrid
  enlarge_feature_map_size: True
  fuse_qkv_proj: False

dataset:
  validate_dataset:
      data_path: ['/data/ocdnet_vit/test']
      args:
        pre_processes:
          - type: Resize2D
            args:
              short_size:
                - 640
                - 640
              resize_text_polys: true
        img_mode: BGR
        filter_keys: []
        ignore_tags: ['*', '###']
      loader:
        batch_size: 1
        shuffle: false
        pin_memory: false
        num_workers: 1

参数	数据类型	默认值	描述	支持的值
`checkpoint`	string		要剪枝的 PyTorch 模型的路径	unix 路径
`ch_sparsity`	float	0.1	剪枝阈值	0.0 ~ 1.0
`results_dir`	string		结果目录的路径	Unix 路径
`round_to`	unsigned int		将通道四舍五入到最接近 round_to 的倍数。例如，round_to=8 表示通道将四舍五入到 8x。	>0
`p`	unsigned int		用于估计通道重要性的范数度。默认值：2	>0
`verbose`	bool		一个标志，指示是否打印剪枝信息，默认值：True	true/false
`fuse_qkv_proj`	bool		一个标志，指示是否融合 qkv 投影，默认值：True，仅当使用 fan-tiny 主干时才需要设置为 True。	true/false

使用以下命令在 OCDNet 模型上运行剪枝。

复制
已复制!

            
            tao model ocdnet prune -e $SPECS_DIR/prune.yaml \
                 prune.checkpoint=$RESULTS_DIR/train/model_best.pth \
                 prune.results_dir=$RESULTS_DIR/prune \
                 [prune.<prune_option>=<prune_option_value>]

必需参数

-e, --experiment_spec_file：用于设置剪枝实验的实验规范文件。

可选参数

prune.<prune_option>：剪枝选项。

剪枝后，剪枝模型可用于重新训练（即微调）。要开始重新训练，您需要将 load_pruned_graph 参数设置为 true，并将 pruned_graph_path 参数设置为指向从剪枝生成的模型。

注意

当重新训练、评估、对具有剪枝结构的模型执行推理或导出模型时，您需要将 load_pruned_graph 设置为 true，以便导入新剪枝的模型结构。有关更多详细信息，请参见示例。

以下是使用剪枝模型运行训练的示例

复制
已复制!

            
            tao model ocdnet train -e $SPECS_DIR/train.yaml \
                 train.results_dir=$RESULTS_DIR/retrain \
                 model.load_pruned_graph=true \
                 model.pruned_graph_path=$RESULTS_DIR/prune/pruned_$ch_sparsity.pth

以下是恢复针对剪枝模型进行训练的示例

复制
已复制!

            
            tao model ocdnet train -e $SPECS_DIR/train.yaml \
                 train.results_dir=$RESULTS_DIR/retrain \
                 model.load_pruned_graph=true \
                 model.pruned_graph_path=$RESULTS_DIR/prune/pruned_$ch_sparsity.pth
                 train.resume_training_checkpoint_path=$RESULTS_DIR/retrain/resume.pth

以下是针对剪枝模型运行评估的示例

复制
已复制!

            
            tao model ocdnet evaluate -e $SPECS_DIR/evaluate.yaml \
                 train.results_dir=$RESULTS_DIR/evaluate \
                 model.load_pruned_graph=true \
                 model.pruned_graph_path=$RESULTS_DIR/prune/pruned_$ch_sparsity.pth
                 evaluate.checkpoint==$RESULTS_DIR/train/model_best.pth

以下是针对剪枝模型运行推理的示例

复制
已复制!

            
            tao model ocdnet inference -e $SPECS_DIR/inference.yaml \
                 model.load_pruned_graph=true \
                 model.pruned_graph_path=$RESULTS_DIR/prune/pruned_$ch_sparsity.pth
                 inference.checkpoint=$RESULTS_DIR/train/model_best.pth \
                 inference.input_folder=$DATA_DIR/test/img \
                 inference.results_dir=$RESULTS_DIR/infer

以下是针对剪枝模型运行导出的示例

复制
已复制!

            
            tao model ocdnet export -e $SPECS_DIR/export.yaml \
                 model.load_pruned_graph=true \
                 model.pruned_graph_path=$RESULTS_DIR/prune/pruned_$ch_sparsity.pth
                 export.checkpoint=$RESULTS_DIR/train/model_best.pth \
                 export.onnx_file=$RESULTS_DIR/export/model_best.onnx

导出模型

export 参数定义导出过程的超参数。

复制
已复制!

            
            model:
  load_pruned_graph: False
  pruned_graph_path: '/results/prune/pruned_0.1.pth'
  backbone: deformable_resnet18
export:
  results_dir: /results/export
  checkpoint: '/results/train/model_best.pth'
  onnx_file: '/results/export/model_best.onnx'
  width: 1280
  height: 736
dataset:
  validate_dataset:
    data_path: ['/data/ocdnet/test']

参数	数据类型	默认值	描述	支持的值
`checkpoint`	string		要导出的 PyTorch 模型的路径	Unix 路径
`onnx_file`	string		ONNX 文件的路径	Unix 路径
`opset_version`	unsigned int	11	导出的 ONNX 的 opset 版本	>0
`input_width`	unsigned int	1280	输入宽度	>0
`input_height`	unsigned int	736	输入高度	>0

复制
已复制!

            
            tao ocdnet export -e $SPECS_DIR/export.yaml
                  export.checkpoint=<path_to_pth_file>
                  export.onnx_file=<path_to_onnx_file>
                  [export.<export_option>=<export_option_value>]

必需参数

-e, --experiment_spec：用于设置导出的实验规范文件
export.checkpoint：要导出的 .pth 模型。
export.onnx_file：保存导出模型的路径

可选参数

export.<export_option>：导出选项。

TensorRT 引擎生成、验证和 INT8 校准

有关部署，请参阅 TAO 部署文档。

注意

如果您未使用 tao deploy 运行 OCDNet TensorRT 引擎，换句话说，如果您在 x86 平台中运行 nm -gDC /usr/lib/x86_64-linux-gnu/libnvinfer_plugin.so |grep ModulatedDeformableConvPlugin 或在 Jetson 平台中运行 nm -gDC /usr/lib/aarch64-linux-gnu/libnvinfer_plugin.so |grep ModulatedDeformableConvPlugin 时没有输出，则需要编译/替换 TensorRT OSS 插件，因为 OCDNet 需要 modulatedDeformConvPlugin。

获取 TensorRT 存储库

复制
已复制!

            
            git clone -b release/8.6 https://github.com/NVIDIA/TensorRT.git
cd TensorRT
git submodule update --init --recursive

编译 TensorRT libnvinfer_plugin.so 文件

复制
已复制!

            
            mkdir build && cd build
# On X86 platform
cmake ..
# On Jetson platform
cmake .. -DTRT_LIB_DIR=/usr/lib/aarch64-linux-gnu/
make nvinfer_plugin -j12

libnvinfer_plugin.so.8.6.x 在 build 文件夹下生成。请注意，x 取决于实际的次版本。

替换默认插件库。请注意，确切的插件名称取决于系统中安装的 TensorRT 版本。

复制
已复制!

            
            # On X86 platform, for example, if the default plugin is /usr/lib/x86_64-linux-gnu/libnvinfer_plugin.so.8.5.2, then
cp libnvinfer_plugin.so.8.6.x /usr/lib/x86_64-linux-gnu/libnvinfer_plugin.so.8.5.2
# On Jetson platform, for example, if the default plugin is /usr/lib/aarch64-linux-gnu/libnvinfer_plugin.so.8.5.2, then
cp libnvinfer_plugin.so.8.6.x /usr/lib/aarch64-linux-gnu/libnvinfer_plugin.so.8.5.2

部署到 DeepStream

有关将 OCDNet 模型部署到 DeepStream 的更多信息，请参阅 nvOCDR 页面。您可以将 nvOCDR 与 DeepStream 示例或 Triton 推理服务器一起运行。具体来说，nvOCDR Triton 可以支持针对高分辨率图像的推理。简而言之，它将在保持宽高比的同时调整图像大小，然后将图像平铺成小块，并运行 OCDNet 以获得输出，然后合并结果。这在模型以较小分辨率训练但将针对较高分辨率图像运行推理的情况下，对于提高 hmean 非常有用。对于分辨率不高图像，您也可以设置 resize_keep_aspect_ratio:true，这对于提高 hmean 非常有用，因为图像在不失真的情况下调整大小。