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

SiameseOI

SiameseOI 是 NVIDIA 开发的用于 PCB 数据的光学检测模型，包含在 TAO 中。SiameseOI 支持以下任务

  • 训练

  • 评估

  • 推理

  • 导出

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

复制
已复制!
            

            
tao model optical_inspection <sub_task> <args_per_subtask>

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

SiameseOI 的数据输入

SiameseOI 要求数据以图像文件夹和 CSV 文件的形式提供。有关 SiameseOI 输入数据格式的更多信息，请参阅数据标注格式页面。

创建训练实验规范文件

配置自定义数据集

本节提供了一个示例配置和命令，用于使用上述数据集格式训练 SiameseOI。您需要根据您的输入数据集配置 augmentation_config 均值和标准差。

这是一个示例规范文件，用于使用数据标注格式在自定义数据集上使用自定义骨干网络训练 SiameseOI 模型。

复制
已复制!
            

            
results_dir: /path/to/experiment_results
model:
  model_type: Siamese_3
  model_backbone: custom
  embedding_vectors: 5
  margin: 2.0
dataset:
  train_dataset:
    csv_path: /path/to/split/train.csv
    images_dir: /path/to/images_dir/
  validation_dataset:
    csv_path: /path/to/split/val.csv
    images_dir: /path/to/images_dir/
  image_ext: .jpg
  batch_size: 32
  workers: 8
  fpratio_sampling: 0.1
  num_input: 4
  input_map:
    LowAngleLight: 0
    SolderLight: 1
    UniformLight: 2
    WhiteLight: 3
  concat_type: linear
  grid_map:
    x: 2
    y: 2
  image_width: 100
  image_height: 100
  augmentation_config:
    rgb_input_mean: [0.485, 0.456, 0.406]
    rgb_input_std: [0.229, 0.224, 0.225]
train:
  optim:
    type: Adam
    lr: 0.0005
  loss: contrastive
  num_epochs: 10
  checkpoint_interval: 5
  validation_interval: 5
  results_dir: "${results_dir}/train"
  seed: 1234

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

训练

参数 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 配置 None 包含 optim 章节中详述的 SiameseOI 优化器的可配置参数。
loss str contrastive 训练期间使用的损失函数

optim

复制
已复制!
            

            
optim:
  lr: 0.0005

参数 Datatype 默认值 描述 支持的值
lr float 0.0005 学习率 >=0.0

模型

以下示例 model 配置提供了更改 SiameseOI 架构以进行训练的选项。

复制
已复制!
            

            
model:
  model_type: Siamese_3
  model_backbone: custom
  embedding_vectors: 5
  margin: 2.0

以下示例 model 在 SiameseOI 评估/推理期间使用。

参数 Datatype 默认值 描述 支持的值
model_type string Siamese_3 来自支持的自定义模型架构的默认模型架构 Siamese_3, Siamese_1
model_backbone string custom 要使用的骨干网络的名称 custom
embedding_vectors int 5 在计算欧几里得距离之前，模型最终输出的嵌入维度
margin float 2.0 确定正负样本对嵌入之间最小距离的阈值参数

数据集

dataset 参数定义数据集源、训练批次大小、数据增强和预处理。下面提供了一个示例 dataset。

复制
已复制!
            

            
dataset:
  train_dataset:
    csv_path: /path/to/split/train.csv
    images_dir: /path/to/images_dir/
  validation_dataset:
    csv_path: /path/to/split/val.csv
    images_dir: /path/to/images_dir/
  image_ext: .jpg
  batch_size: 32
  workers: 8
  fpratio_sampling: 0.1
  num_input: 4
  input_map:
    LowAngleLight: 0
    SolderLight: 1
    UniformLight: 2
    WhiteLight: 3
  concat_type: linear
  grid_map:
    x: 2
    y: 2
  image_width: 100
  image_height: 100
  augmentation_config:
    rgb_input_mean: [0.485, 0.456, 0.406]
    rgb_input_std: [0.229, 0.224, 0.225]

参数 Datatype 默认值 描述 支持的值
train_dataset Dict – 训练数据集的图像目录和 CSV 文件路径
validation_dataset Dict – 验证数据集的图像目录和 CSV 文件路径
image_ext str .jpg 数据集中图像的文件扩展名 string
batch_size int 32 每个批次的样本数 string
workers int 8 用于数据加载的工作进程数
fpratio_sampling int 0.1 假阳性样本的采样比率 >0
num_input int 4 每个输入图像的光照条件数* >0
input_map Dict – 将光照条件映射到指定连接顺序的索引*
concat_type string linear 用于不同图像光照条件的连接类型 linear, grid

grid_map

Dict

Dict

dict 配置

None

None

None

用于定义网格尺寸以将图像连接为网格的参数

* x：沿 x 轴的图像数量

* y：沿 y 轴的图像数量

Dict
input_width int 100 输入图像的宽度 >0
input_height int 100 输入图像的高度 >0

augmentation_config

Dict

List[float]

List[float]

None

[0.485, 0.456, 0.406]

[0.229, 0.224, 0.225]

图像归一化配置，其中包含以下参数

* rgb_input_mean：预处理要减去的均值

* rgb_input_std：图像要除以的标准差

>=0.0

>=0.0

* 有关指定光照条件的更多信息，请参阅数据集标注格式定义。

训练模型

使用以下命令运行 SiameseOI 训练

复制
已复制!
            

            
tao model optical_inspection train [-h] -e <experiment_spec>
                             [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：用于设置训练实验的实验规范文件

可选参数

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

注意

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

最新的检查点也会保存为 oi_model_latest.pth。如果 oi_model_latest.pth 存在于 train.results_dir 中，则训练会自动从其恢复。如果提供了 train.resume_training_checkpoint_path，则它会被覆盖。

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

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

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

创建测试实验规范文件

这是一个示例规范文件，用于测试训练后的 SiameseOI 模型的评估和推理。

复制
已复制!
            

            
results_dir: /path/to/experiment_results
model:
  model_type: Siamese_3
  model_backbone: custom
  embedding_vectors: 5
  margin: 2.0
dataset:
  validation_dataset:
    csv_path: /path/to/split/val.csv
    images_dir: /path/to/images_dir/
  image_ext: .jpg
  batch_size: 32
  workers: 8
  num_input: 4
  input_map:
    LowAngleLight: 0
    SolderLight: 1
    UniformLight: 2
    WhiteLight: 3
  concat_type: linear
  grid_map:
    x: 2
    y: 2
  image_width: 100
  image_height: 100
  augmentation_config:
    rgb_input_mean: [0.485, 0.456, 0.406]
    rgb_input_std: [0.229, 0.224, 0.225]
evaluate:
  num_gpus: 1
  gpu_ids: [0]
  checkpoint: "${results_dir}/train/oi_model_lastest.pth"
  results_dir: "${results_dir}/evaluate"
inference:
  num_gpus: 1
  gpu_ids: [0]
  checkpoint: "${results_dir}/train/oi_model_latest.pth"
  results_dir: "${results_dir}/inference"

评估模型

使用以下命令运行 SiameseOI 评估

复制
已复制!
            

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

光学检测目前不支持多 GPU 评估。

必需参数

  • -e, --experiment_spec：用于设置评估实验的实验规范文件。

  • evaluate.checkpoint：要评估的 .pth 模型。

可选参数

在模型上运行推理

使用以下命令对 SiameseOI 运行推理，使用 .tlt 模型

复制
已复制!
            

            
tao model optical_inspection inference [-h] -e <experiment spec file>
                             inference.checkpoint=<model to be inferenced>
                             [inference.<inference_option>=<inference_option_value>]
                             [inference.gpu_ids=<gpu indices>]
                             [inference.num_gpus=<number of gpus>]

必需参数

  • -e, --experiment_spec：用于设置推理实验的实验规范文件。

  • inference.checkpoint：要推理的 .pth 模型。

可选参数

导出模型

这是一个示例规范文件，用于导出训练后的 SiameseOI 模型

复制
已复制!
            

            
export:
  checkpoint: "${results_dir}/train/oi_model_epoch=004.pth"
  results_dir: "${results_dir}/export"
  onnx_file: "${export.results_dir}/oi_model.onnx"
  batch_size: 32

使用以下命令导出模型

复制
已复制!
            

            
tao model optical_inspection export [-h] -e <experiment spec file>
                             export.checkpoint=<model to export>
                             export.onnx_file=<onnx path>
                             [export.<export_option>=<export_option_value>]

必需参数

  • -e, --experiment_spec：实验规范文件的路径。

  • export.checkpoint：要导出的 .pth 模型。

  • export.onnx_file：保存 .etlt 或 .onnx 模型的路径。

可选参数

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

有关部署，请参阅适用于 SiameseOI 的 TAO 部署文档。
上一篇 光学检测
下一篇 姿态分类
© 版权所有 2024, NVIDIA。 上次更新于 2024 年 10 月 15 日。
此处内容