TAO v5.5.0
NVIDIA 文档中心  NVIDIA TAO  TAO v5.5.0  OCRNet

OCRNet

OCRNet 是一个用于识别图像中字符的模型。它支持以下任务

  • dataset_convert

  • train

  • evaluate

  • prune

  • inference

  • export

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

复制
已复制！
            

            
tao model ocrnet <sub_task> <args_per_subtask>

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

准备数据集

OCRNet 的训练数据集和评估数据集为 LMDB 格式。您可以使用 dataset_convert 将原始图像和标签转换为 LMDB 格式。原始数据集应按以下结构组织

复制
已复制！
            

            
/Dataset
    /images
        0000.jpg
        0001.jpg
        0002.jpg
        ...
    gt_list.txt
    characters_list.txt

gt_list.txt 文件包含图像的所有真实文本，每张图像及其对应的标签用一行文本指定

复制
已复制！
            

            
0000.jpg abc
0001.jpg defg
0002.jpg zxv
...

存在一个 characters_list.txt 文件，其中包含数据集中找到的所有字符。每个字符占一行。

创建实验规范文件

实验规范文件包括 OCRNet 支持的所有任务（train/evaluate/inference/prune/export）的参数。以下是 OCRNet get_started notebook 中使用的示例规范文件

复制
已复制！
            

            
results_dir: /results
encryption_key: nvidia_tao
model:
  TPS: True
  backbone: ResNet
  feature_channel: 512
  sequence: BiLSTM
  hidden_size: 256
  prediction: CTC
  quantize: False
  input_width: 100
  input_height: 32
  input_channel: 1
dataset:
  train_dataset_dir: []
  val_dataset_dir: /data/test/lmdb
  character_list_file: /data/character_list
  max_label_length: 25
  batch_size: 32
  workers: 4
  augmentation:
    keep_aspect_ratio: False
train:
  seed: 1111
  gpu_ids: [0]
  optim:
    name: "adadelta"
    lr: 1.0
  clip_grad_norm: 5.0
  num_epochs: 10
  checkpoint_interval: 2
  validation_interval: 1
evaluate:
  gpu_id: 0
  checkpoint: "??"
  test_dataset_dir: "??"
  results_dir: "${results_dir}/evaluate"
prune:
  gpu_id: 0
  checkpoint: "??"
  results_dir: "${results_dir}/prune"
  prune_setting:
    mode: experimental_hybrid
    amount: 0.4
    granularity: 8
    raw_prune_score: L1
inference:
  gpu_id: 0
  checkpoint: "??"
  inference_dataset_dir: "??"
  results_dir: "${results_dir}/inference"
export:
  gpu_id: 0
  checkpoint: "??"
  results_dir: "${results_dir}/export"
dataset_convert:
  input_img_dir: "??"
  gt_file: "??"
  results_dir: "${results_dir}/convert_dataset"
gen_trt_engine:
  onnx_file: "??"
  results_dir: "${results_dir}/convert_dataset"

参数 数据类型 默认值 描述 支持的值
model dict config – 模型架构的配置
dataset dict config – 数据集的配置
train dict config – 训练任务的配置
evaluate dict config – 评估任务的配置
inference dict config – 推理任务的配置
encryption_key string None 用于加密和解密模型文件的加密密钥
results_dir string /results 实验结果保存的目录
prune dict config – 剪枝的配置
export dict config – 导出的配置
dataset_convert dict config – 数据集转换的配置

model

model 参数提供更改 OCRNet 架构的选项。

复制
已复制！
            

            
model:
  TPS: True
  backbone: ResNet
  feature_channel: 512
  sequence: BiLSTM
  hidden_size: 256
  prediction: CTC
  quantize: False

参数

Datatype

默认值

描述

支持的值
TPS Boolean False 一个标志，用于为 OCRNet 输入启用薄板样条插值 True/False
num_fiducial Unsigned int 20 TPS 的特征点数量 >4
backbone String ResNet OCRNet 模型的主干网络 ResNet, ResNet2X, FAN_tiny_2X
feature_channel Unsigned int 512 主干网络输出特征的通道数 >0
sequence String BiLSTM OCRNet 模型的序列模块 BiLSTM
hidden_size Unsigned int 256 BiLSTM 隐藏层的通道数 >0
prediction String CTC 用于编码和解码输出特征的方法 CTC, Attn
input_width Unsigned int 100 输入图像宽度 >4
input_height Unsigned int 32 输入图像高度 >32
input_channel Unsigned int 1 输入图像通道 1,3
quantize Boolean False 一个标志，用于在 OCRNet 主干网络中启用量化和反量化节点 True/False

dataset

dataset 参数提供设置训练和评估中使用的数据集的选项。

复制
已复制！
            

            
dataset:
  train_dataset_dir: [/data/train/lmdb]
  val_dataset_dir: /data/test/lmdb
  character_list_file: /data/character_list
  max_label_length: 25
  batch_size: 32
  workers: 4
  augmentation:
    keep_aspect_ratio: False
    aug_prob: 0.3
    reverse_color_prob: 0.5
    rotate_prob: 0.5
    max_rotation_degree: 5
    blur_prob: 0.5
    gaussian_radius_list: [1, 2, 3, 4]

参数

Datatype

默认值

描述

支持的值
train_dataset_dir List of String None 训练数据集的绝对路径列表。目前，仅支持长度为 1 的列表。 List of String
val_dataset_dir String None 评估数据集的绝对路径 数据集绝对路径
character_list_file String None 字符列表文件的绝对路径 绝对文件路径
max_label_length Unsigned int 25 真实标签的最大长度 >0
batch_size Unsigned int 32 训练的批次大小 >0
workers Unsigned int 4 并行预处理训练数据的工作线程数 >=0
augmentation Dict config – 数据增强配置。 –

augmentation

augmentation 参数提供在训练期间设置数据增强管道的选项。

复制
已复制！
            

            
augmentation:
  keep_aspect_ratio: False
  aug_prob: 0.3
  reverse_color_prob: 0.5
  rotate_prob: 0.5
  max_rotation_degree: 5
  blur_prob: 0.5
  gaussian_radius_list: [1, 2, 3, 4]

参数

Datatype

默认值

描述

支持的值
keep_aspect_ratio Bool False 一个标志，用于在将图像大小调整为模型输入大小时保持宽高比 False/True
aug_prob Float 0.0 对输入图像应用以下数据增强的概率 [0, 1]
reverse_color_prob Float 0.5 反转输入图像颜色的概率 [0, 1]
rotate_prob Float 0.5 随机旋转输入图像的概率 [0, 1]
max_rotation_degree Float 0.5 图像将旋转的最大角度 >=0
blur_prob Float 0.5 模糊输入图像的概率 [0, 1]
gaussian_radius_list List of integer [1, 2, 3, 4] 在图像上应用高斯模糊时使用的半径列表 –

train

train 参数提供设置训练超参数的选项。

复制
已复制！
            

            
train:
  seed: 1111
  gpu_ids: [0]
  optim:
    name: "adadelta"
    lr: 1.0
  clip_grad_norm: 5.0
  num_epochs: 10
  checkpoint_interval: 2
  validation_interval: 1

参数 Datatype 默认值 描述 支持的值
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 用于保存训练结果的目录
optim Dict config – 优化器的配置 –
clip_grad_norm Float 5.0 要裁剪的梯度 L2 范数幅度的阈值 >4
distributed_strategy String ddp 用于多 GPU 训练的分布式策略 ddp
pretrained_model_path String None 预训练权重的绝对路径 –
quantize_model_path String None 用于量化感知训练的预训练模型的绝对路径 –
model_ema Bool False 在训练中启用模型指数移动平均 False/True

optim

optim 提供设置训练优化器的选项。

复制
已复制！
            

            
optim:
  name: "adadelta"
  lr: 1.0

参数

Datatype

默认值

描述

支持的值
name String adadelta 优化器类型 adadelta, adam
lr Float 1.0 训练的初始学习率 >0.0

evaluate

evaluate 参数提供设置评估超参数的选项。

复制
已复制！
            

            
evaluate:
  checkpoint: "??"
  test_dataset_dir: "??"
  results_dir: "${results_dir}/evaluate"

参数

Datatype

默认值

描述

支持的值
checkpoint String – 用于评估的模型检查点的绝对路径 –
results_dir String /results/evaluate 用于保存评估结果的目录
num_gpus Unsigned int 1 用于分布式评估的 GPU 数量 >0
gpu_ids List[int] [0] 用于分布式评估的 GPU 索引列表
test_dataset_dir String – 评估 LMDB 数据集的绝对路径 –
batch_size Unsigned int 1 评估批次大小 >0

prune

prune 参数提供设置剪枝超参数的选项。

复制
已复制！
            

            
gpu_id: 0
checkpoint: "??"
results_dir: "${results_dir}/prune"
prune_setting:
  mode: experimental_hybrid
  amount: 0.4
  granularity: 8
  raw_prune_score: L1

参数

Datatype

默认值

描述

支持的值
checkpoint String – 用于剪枝的模型检查点的绝对路径 –
gpu_id Unsigned int 0 GPU 设备索引 有效的 GPU 索引
results_dir String – 剪枝日志的绝对路径 –
pruned_file String – 用于存储剪枝模型检查点的绝对路径 –
prune_setting Dict config – 剪枝超参数 –

prune_setting

prune_setting 参数包含剪枝算法的选项

参数

Datatype

默认值

描述

支持的值
mode String amount 要剪枝的模型检查点的绝对路径
  • amount: 根据重要性剪枝权重数量比率
  • threshold: 剪枝重要性小于阈值的权重
  • experimental_hybrid: 使用 threshold 和 amount 的混合方法剪枝权重
 amount, threshold, experimental_hybrid
amount Float – amount 和 experimental_hybrid 模式的数量值 [0, 1]
threshold Float – threshold 模式的阈值 >=0
granularity Unsigned int 8 剪枝层的粒度。剪枝层输出通道的数量将是粒度的倍数。 >0
raw_prune_score Dict config L1 用于计算权重重要性的方法 L1, L2

inference

inference 参数提供推理选项。

复制
已复制！
            

            
inference:
  checkpoint: "??"
  inference_dataset_dir: "??"
  results_dir: "${results_dir}/inference"

参数

Datatype

默认值

描述

支持的值
checkpoint String – 用于推理的模型检查点的绝对路径 –
results_dir String /results/inference 用于保存推理结果的目录
num_gpus Unsigned int 1 用于分布式推理的 GPU 数量 >0
gpu_ids List[int] [0] 用于分布式推理的 GPU 索引列表
inference_dataset_dir String – 推理图像目录的绝对路径 –
batch_size Unsigned int 1 推理批次大小 >0

export

export 参数提供导出选项。

复制
已复制！
            

            
export:
  gpu_id: 0
  checkpoint: "??"
  results_dir: "${results_dir}/export"

参数

Datatype

默认值

描述

支持的值
checkpoint String – 用于导出的模型检查点的绝对路径 –
gpu_id Unsigned int 0 GPU 设备索引 有效的 GPU 索引
onnx_file String – 导出 ONNX 文件的绝对路径 –
results_dir String – 导出输出的绝对路径 –

dataset_convert

dataset_convert 参数提供设置数据集转换的选项。

复制
已复制！
            

            
dataset_convert:
  input_img_dir: "??"
  gt_file: "??"
  results_dir: "${results_dir}/convert_dataset"

参数

Datatype

默认值

描述

支持的值
input_img_dir String – 图像目录的绝对路径 –
gt_file String – 真实标签文件的绝对路径 –
results_dir String – dataset_convert 的绝对路径（即 LMDB 数据集和日志） –

转换数据集

使用以下命令将原始数据集转换为 LMDB 格式

复制
已复制！
            

            
tao model ocrnet dataset_convert -e <experiment_spec_file>
                 [results_dir=<global_results_dir>]
                 [dataset_convert.<dataset_convert_option>=<dataset_convert_value>]

必需参数

  • -e, --experiment_spec_file: 实验规范文件的路径。

可选参数

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

  • results_dir: 全局结果目录。dataset_convert 结果将保存在 results_dir/dataset_convert 中。

  • dataset_convert.<dataset_convert_option>: dataset_convert 选项。

训练模型

使用以下命令启动 OCRNet 训练

复制
已复制！
            

            
tao model ocrnet 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.optim.<optim_option>=<optim_option_value>]
                 [train.gpu_ids=<gpu indices>]
                 [train.num_gpus=<number of gpus>]

必需参数

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

可选参数

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

注意

对于训练、评估和推理，我们为每个相应的任务公开 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'

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

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

  • 指定一个新的空结果目录（推荐）

  • 从结果目录中删除最新的检查点

评估模型

使用以下命令启动 OCRNet 评估

复制
已复制！
            

            
tao model ocrnet evaluate -e <experiment_spec_file>
                 evaluate.checkpoint=<model to be evaluated>
                 [model.<model_option>=<model_option_value>]
                 [dataset.<dataset_option>=<dataset_option_value>]
                 [evaluate.<evaluate_option>=<evaluate_option_value>]
                 [evaluate.gpu_ids=<gpu indices>]
                 [evaluate.num_gpus=<number of gpus>]

OCRNet 目前不支持多 GPU 评估。

必需参数

  • -e, --experiment_spec_file: 用于设置评估实验的实验规范文件。这应该与训练规范文件相同。

  • evaluate.checkpoint: .pth 模型。

可选参数

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

OCRNet 目前不支持多 GPU 评估。

剪枝模型

使用以下命令启动 OCRNet 剪枝

复制
已复制！
            

            
tao model ocrnet prune -e <experiment_spec_file>
                 [results_dir=<global_results_dir>]
                 [model.<model_option>=<model_option_value>]
                 [dataset.<dataset_option>=<dataset_option_value>]
                 [prune.<prune_option>=<prune_option_value>]
                 [prune.prune_setting.<prune_setting_option>=<prune_setting_value>]

必需参数

  • -e, --experiment_spec_file: 实验规范文件的路径。

可选参数

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

注意

如果在剪枝图上运行训练、评估或推理，则在运行相应任务时，必须提供 model.pruned_graph_path 参数。它应与剪枝任务中为 prune.pruned_file 提供的值相同。

使用模型进行推理

使用以下命令启动 OCRNet 推理

复制
已复制！
            

            
tao model ocrnet inference -e <experiment_spec_file>
                 [results_dir=<global_results_dir>]
                 [model.<model_option>=<model_option_value>]
                 [dataset.<dataset_option>=<dataset_option_value>]
                 [inference.<inference_option>=<inference_option_value>]

OCRNet 目前不支持多 GPU 推理。

必需参数

  • -e, --experiment_spec_file: 实验规范文件的路径。

可选参数

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

  • results_dir: 全局结果目录。推理结果将保存在 results_dir/inference 中。

  • model.<model_option>: 模型选项。

  • dataset.<dataset_option>: 数据集选项。

  • inference.<inference_option>: 推理选项。

导出模型

使用以下命令将 OCRNet PyTorch 检查点导出到 ONNX 模型

复制
已复制！
            

            
tao model ocrnet export -e <experiment_spec_file>
                 [results_dir=<global_results_dir>]
                 [model.<model_option>=<model_option_value>]
                 [dataset.<dataset_option>=<dataset_option_value>]
                 [export.<export_option>=<export_option_value>]

必需参数

  • -e, --experiment_spec_file: 实验规范文件的路径。

可选参数

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

  • results_dir: 全局结果目录。导出结果将保存在 results_dir/export 中。

  • model.<model_option>: 模型选项。

  • dataset.<dataset_option>: 数据集选项。

  • export.<export_option>: 导出选项。

TensorRT 引擎生成和验证

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

部署到 DeepStream

有关 DeepStream 集成，请参阅 将 nvOCDR 部署到 DeepStream。
上一个 字符识别
下一个 VisualChangeNet
© 版权所有 2024, NVIDIA。 上次更新时间：2024 年 10 月 15 日。
此处内容