tao-train-optical-inspection

光学检测

独立安装？ 如果此会话并非由 TAO skill bank 插件初始化，请先运行 tao-setup 技能（主机预检、凭据、跨技能发现）。

使用 Siamese 网络进行缺陷检测的光学检测。比较成对图像，以检测制造缺陷、异常或质量问题。

为预训练的 Siamese 权重设置 train.pretrained_model_path。

对于 TAO Deploy TensorRT 操作（gen_trt_engine、TensorRT evaluate 和 TensorRT inference），请先阅读 references/tao-deploy-optical-inspection.md。父级 PyT 容器不提供 optical_inspection gen_trt_engine；TensorRT 引擎只能通过 Deploy 生成。Deploy 规范模板位于此技能的 references/ 文件夹中，文件名前缀为 spec_template_deploy_*.yaml。

Dataclass Schema

生成的 TAO Core schema 打包在 schemas/<action>.schema.json 中，schemas/manifest.json 则列出可用操作。每个生成的 schema 还会根据 schema 顶层的 default 字段生成 references/spec_template_<action>.yaml。AutoML 启用状态通过 references/skill_info.yaml 中模型层级的 automl_enabled 声明。某项操作要能够运行 AutoML，必须存在且可解析 schemas/<action>.schema.json 和 references/spec_template_<action>.yaml。请使用打包的所选操作 schema 获取 automl_default_parameters、automl_disabled_parameters、默认值、最小/最大边界、枚举值、选项权重、数学条件、依赖项和常用参数。不要假设运行时存在 ~/tao-core；维护者会在打包 skill bank 前重新生成 schema 和模板。

Train 操作策略

此模型在模型层级启用了 AutoML。处理任何训练阶段请求前，请阅读 references/skill_info.yaml，并根据显式的 automl_policy 值或用户的工作流请求确定本次运行的覆盖设置。默认使用 automl_policy: on，且在新的启动提示中只提供 on / off 选项。将“关闭 AutoML”“禁用 AutoML”“不使用 HPO”或“普通训练”等表述视为仅对本次运行设置 automl_policy: off。当 automl_policy: on、automl_enabled: true，且已打包 schemas/train.schema.json 和 references/spec_template_train.yaml 时，默认通过 tao-skill-bank:tao-run-automl 执行 train 操作，并传入此模型的 skill_dir。保留工作流/应用对数据集、规范、输出目录、GPU/平台设置、父级检查点和 automl_policy 的覆盖设置。只有在 automl_policy: off 或缺少已打包的 train schema/模板时，才直接进行模型训练；若缺少 schema，请说明此模型已启用 AutoML，但在生成 schema 前无法运行 AutoML。

evaluate、inference、export 和 Deploy 流程等非训练操作仍由此模型技能处理。每次运行的 automl_policy 覆盖设置不会更改模型元数据。

训练要求

• 数据集类型： optical_inspection
• 格式： 默认
• AutoML 训练指标： val_acc，其中 direction=maximize
• 独立评估指标： test_acc，其中 direction=maximize

各操作的数据集要求

操作	规范键	来源	文件	是否为列表？
评估	dataset.test_dataset.images_dir	eval_dataset	images.tar.gz	否
评估	dataset.test_dataset.csv_path	eval_dataset	dataset.csv	否
推理	dataset.infer_dataset.images_dir	inference_dataset	images.tar.gz	否
推理	dataset.infer_dataset.csv_path	inference_dataset	dataset.csv	否
训练	dataset.train_dataset.images_dir	train_datasets	images.tar.gz	否
训练	dataset.train_dataset.csv_path	train_datasets	dataset.csv	否
训练	dataset.validation_dataset.images_dir	eval_dataset	images.tar.gz	否
训练	dataset.validation_dataset.csv_path	eval_dataset	dataset.csv	否
训练	dataset.test_dataset.images_dir	eval_dataset	images.tar.gz	否
训练	dataset.test_dataset.csv_path	eval_dataset	dataset.csv	否

images.tar.gz 是传输文件。完成暂存／解压后， dataset.*_dataset.images_dir 必须指向内层目录，该目录直接包含 golden/ 和 CSV 中使用的棋盘目录。

dataset.csv 约定

Optical Inspection 加载器按列名读取 CSV，并按以下方式构造每组孪生比较的两侧：

<images_dir>/<input_path>/<object_name>_<lighting><image_ext>
<images_dir>/<golden_path>/<object_name>_<lighting><image_ext>

加载器也接受绝对路径形式的 input_path 和 golden_path 值；绝对路径会替代 images_dir。建议使用相对目录，因为暂存后仍可移植。两个路径列中都不要填写文件名，也不要在 object_name 中填写光照后缀或扩展名。

列	类型	必填	允许的值和含义
input_path	字符串目录路径	是	板卡／采集组件目录，相对于 images_dir 或使用绝对路径。
golden_path	字符串目录路径	是	黄金样本／参考组件目录，相对于 images_dir 或使用绝对路径。
label	字符串	是	大小写完全匹配的 PASS 表示无缺陷（类别 0）。任何其他非空缺陷名称都表示有缺陷（类别 1），例如 missing、shift、excess_solder、lifted_lead、polarity、tombstone 或 upside down。pass 无效，因为加载器会将其静默地视为缺陷。
object_name	字符串文件名主干	是	组件标识符，例如 C1018@1；不能包含目录、_SolderLight 或 .jpg。加载器会强制将此列设为字符串，因此看起来像数字的标识符会保留其文本形式。

允许包含其他元数据列，但此加载器会忽略它们。值不能为空。以下是一个基于生产目录结构的完整行：

input_path,golden_path,label,object_name
690-5G190-0510-001P1/AOI_B/FXLH_690-5G190-0510-001P1_30332_P_AOI_B_20230317130332/PerComponent,golden/images/690-5G190-0510-001P1BOT/,PASS,C1018@1

使用标准的四光源配置时，每行需要八个文件：

<images_dir>/
├── golden/images/690-2G133-0210-000BOT/
│   ├── R821@1_LowAngleLight.jpg
│   ├── R821@1_SolderLight.jpg
│   ├── R821@1_UniformLight.jpg
│   └── R821@1_WhiteLight.jpg
└── 690-2G133-0210-000/AOI_B/<capture>/PerComponent/
    ├── R821@1_LowAngleLight.jpg
    ├── R821@1_SolderLight.jpg
    ├── R821@1_UniformLight.jpg
    └── R821@1_WhiteLight.jpg

以下三个选择字段必须相互匹配：

• input_map 的键是精确的文件名光照后缀。当前加载器会按 YAML 键的插入顺序遍历，不会按整数值排序。请保持值连续且顺序匹配（0..N-1）。
• num_input 必须等于 input_map 的条目数。加载器会打开每个键，因此不匹配的值不会限制文件列表，并会导致张量／导出形状不匹配。
• concat_type: linear 会按键顺序沿图像高度堆叠输入。四张 128×128 的输入会生成 512×128 的张量；grid_map 会被忽略。
• concat_type: grid 要求 num_input 为偶数，且 grid_map.x * grid_map.y == num_input。按键顺序以行优先方式放置。对于标准的 2 x 2 映射：LowAngle 位于左上角，Solder 位于右上角，Uniform 位于左下角，White 位于右下角。

如果数据集确实只包含 *_SolderLight.jpg，请使用 num_input: 1、input_map: {SolderLight: 0} 和 concat_type: linear。缺少三个变体时，不要声明四个输入。

在训练、评估或推理前，使用与规格相同的数据集设置运行打包的预检程序：

python3 skills/models/tao-train-optical-inspection/scripts/validate_dataset.py \
  --csv /data/optical-inspection/train/dataset.csv \
  --images-dir /data/optical-inspection/train/images \
  --num-input 4 \
  --concat-type grid \
  --grid-x 2 --grid-y 2

对于自定义映射，请按 YAML 键顺序重复使用 --input-map LIGHT=INDEX，并传入规格中的 --image-ext。验证器会报告缺失的列、不安全的 PASS 大小写、无法解析的行目录，以及每个缺失的光照文件及其 CSV 行号。tests/fixtures/dataset/valid/ 下提供了一个包含两行路径占位数据的夹具；它用于验证约定，不是 PCB 训练数据。

在此工具库中找不到已发布的 Optical Inspection 示例数据集。请向数据集所有者获取适用于你的产品或组织的转换后 AOI 数据，然后将其暂存到挂载至容器的路径中。不要使用已提交的验证器夹具进行模型训练。

典型规格覆盖项

每个操作都必须覆盖数据源——Agent 必须根据上方的“每项操作的数据集要求”表构建数据源路径，并将其包含在 spec_overrides 中。

TRAIN_ROOT = "/data/optical-inspection/train"
EVAL_ROOT = "/data/optical-inspection/eval"
INFERENCE_ROOT = "/data/optical-inspection/inference"

这些是从数据所有者处获取并暂存数据后，容器内的示例挂载点，并非下载位置。

train（必需数据源）：

{
    "train.num_epochs": 30,
    "train.checkpoint_interval": 10,
    "train.validation_interval": 10,
    "train.num_gpus": 1,
    "dataset.batch_size": 8,
    "dataset.train_dataset.images_dir": f"{TRAIN_ROOT}/images",
    "dataset.train_dataset.csv_path": f"{TRAIN_ROOT}/dataset.csv",
    "dataset.validation_dataset.images_dir": f"{EVAL_ROOT}/images",
    "dataset.validation_dataset.csv_path": f"{EVAL_ROOT}/dataset.csv",
    "dataset.test_dataset.images_dir": f"{EVAL_ROOT}/images",
    "dataset.test_dataset.csv_path": f"{EVAL_ROOT}/dataset.csv",
}

evaluate（必需数据源）：

{
    "evaluate.checkpoint": "<selected train/AutoML checkpoint>",
    "dataset.test_dataset.images_dir": f"{EVAL_ROOT}/images",
    "dataset.test_dataset.csv_path": f"{EVAL_ROOT}/dataset.csv",
}

对于下游操作，请使用工作流的检查点解析器，而不是猜测文件名。在 Optical Inspection 冒烟测试中，AutoML 可能会生成 model_epoch_000_step_00006.pth；随后恢复训练可能会生成 model_epoch_001_step_00012.pth。最佳检查点操作应使用 AutoML 最佳子作业选定的检查点；指定 epoch 的操作应传入所请求的确切 epoch/step 检查点；只有明确请求“latest”时，才应解析为最新检查点。

export：

{
    "export.checkpoint": "<selected train/AutoML checkpoint>",
    "export.onnx_file": "/results/optical_inspection.onnx",
    "export.input_width": 128,
    "export.input_height": 512,
    "export.batch_size": 1,
}

inference（必需数据源）：

{
    "inference.checkpoint": "<selected train/AutoML checkpoint>",
    "dataset.infer_dataset.images_dir": f"{INFERENCE_ROOT}/images",
    "dataset.infer_dataset.csv_path": f"{INFERENCE_ROOT}/dataset.csv",
}

数据集转换

对于 Optical Inspection，数据集转换是可选的。如果数据集已经是可供 TAO 使用的 Optical Inspection 格式，请直接从 images.tar.gz 和 dataset.csv 划分文件开始，并使用转换后的数据运行 train、evaluate、inference 以及下游检查点、导出和部署操作。

PyT 容器提供 optical_inspection dataset_convert，但此模型技能包不包含 dataset_convert 操作/模板。该转换器需要原始 Factory PCB 布局（root_dataset_dir、train/val/all PCB 目录、golden_csv_dir、project_name 和 bot_top）。数据所有者也可以提供预转换的 Optical Inspection images.tar.gz 和 dataset.csv 划分文件，而无需提供原始 PCB/golden CSV 源数据。不要伪造 PCB 数据集。如果只有转换后的数据可用，请在模型验证报告中将数据集转换标记为 not run: preconverted dataset provided，而非失败或受阻。

在本地使用预转换的传输归档时，请先验证解压后的目录，再编写规格。归档可能会解压出一个 images/ 包装目录；请将 dataset.*.images_dir 指向其中包含 golden/ 以及 dataset.csv 引用的板卡/图像文件夹的内层目录，例如 .../<split>/images/images，而不是外层包装目录。

产品方面仍需后续跟进：发布获得许可、可用于训练的示例数据集，并将现有容器 dataset_convert 入口打包为带有独立 schema/template 的技能操作。本文档和预检变更均未实现这两项工作。

评估数据集

可选。评估数据集使用相同格式（图像 + CSV）。

重要参数

• model.model_type：Siamese 变体。选项包括 Siamese、Siamese_3。
• model.model_backbone：默认值为 custom。
• model.embedding_vectors：嵌入维度数量。默认值为 5。
• train.optim.lr：学习率。默认值为 5e-4。
• dataset.batch_size：训练批次大小。必须大于 1；最小冒烟测试请使用 2 或更高值。
• dataset.num_input：每次比较的输入图像数量。
• dataset.input_map：输入通道/图像对的映射。

多 GPU / 多节点

**启动方式：**由 Lightning 管理（单个 python 进程，由 Lightning 启动工作进程）。

Spec Key	描述	默认值
train.num_gpus	GPU 数量	1
train.gpu_ids	GPU 设备索引	[0]

• 策略：auto（Lightning 会自动选择最佳策略）
• 无需显式配置 num_nodes 或 distributed_strategy——仅支持单节点
• Siamese 网络较轻量，通常单个 GPU 就足够

硬件

最低需要 1 个 GPU，建议使用 1 个 GPU。每个 GPU 需配备 8GB 或更多显存。用于检测的 Siamese 网络较轻量，单个 GPU 即可满足需求。

错误模式

CSV 格式错误：要求包含 input_path,golden_path,label,object_name，然后在启动前，使用规格文件中的 images_dir、光照映射、num_input、拼接设置、网格设置和图像扩展名运行 scripts/validate_dataset.py。

提取后的图像根目录不匹配：如果训练、评估或推理无法找到 dataset.csv 中的路径，请检查提取后的 images.tar.gz 目录结构。符合 TAO 要求的根目录必须包含 golden/ 以及 CSV 中引用的板卡文件夹。对于传输的归档文件，这些内容可能位于提取目标目录的下一层，例如 images/images。

训练批次大小断言失败：Optical Inspection 数据加载器不允许训练时使用 dataset.batch_size: 1。常规运行请保留模板默认值 8，或将 dataset.batch_size: 2 用于最小化 AutoML 冒烟验证。

下游操作加载 PyTorch 检查点失败：对于由同一受信任 TAO 训练／AutoML 工作流生成的检查点，如果当前 PyTorch 默认设置阻止加载完整检查点，请在评估、推理、导出和恢复任务中设置 TORCH_FORCE_NO_WEIGHTS_ONLY_LOAD=1。请勿将此环境变量用于不受信任的检查点。

规格参数／父模型推理

特定模型的推理映射应放在此 MD 文件中，而不是 config.json 中。生成的运行器应读取本节内容，并在调用 create_job() 之前使用 SDK 辅助函数应用这些映射。这与旧版微服务的 infer_params.py 流程一致。

TAO Core optical_inspection.config.json 中的推理映射：

操作	规格字段	推理函数	含义
评估	encryption_key	key	加密密钥
评估	evaluate.checkpoint	parent_model	根据父任务结果文件夹推断出的模型文件
评估	results_dir	output_dir	当前任务的结果目录
导出	encryption_key	key	加密密钥
导出	export.checkpoint	parent_model	根据父任务结果文件夹推断出的模型文件
导出	export.onnx_file	create_onnx_file	输出 ONNX 路径
导出	results_dir	output_dir	当前任务的结果目录
推理	encryption_key	key	加密密钥
推理	inference.checkpoint	parent_model	根据父任务结果文件夹推断出的模型文件
推理	inference.trt_engine	parent_model	根据父任务结果文件夹推断出的模型文件
推理	results_dir	output_dir	当前任务的结果目录
训练	encryption_key	key	加密密钥
训练	results_dir	output_dir	当前任务的结果目录
训练	train.pretrained_model_path	ptm_if_no_resume_model	不存在恢复检查点时使用的 PTM
训练	train.resume_training_checkpoint_path	resume_model	根据当前任务结果文件夹推断出的模型文件

对于 parent_model 或 parent_model_folder，请将上游训练／导出／AutoML 子任务 ID 作为 parent_job_id 传入。SDK 会列出父任务的结果文件夹、筛选检查点产物，并返回选定的模型文件或文件夹。请勿将这些映射重新添加到 config.json，也不要修改生成的运行器脚本来猜测检查点路径。

部署

• tao-deploy-optical-inspection