SiameseOI is an NVIDIA-developed optical inspection model for PCB data and is included in the TAO Toolkit. SiameseOI supports the following tasks:
trainevaluateinferenceexport
These tasks can be invoked from the TAO Toolkit Launcher using the following convention on the command-line:
tao model optical_inspection <sub_task> <args_per_subtask>
Where args_per_subtask are the command-line arguments required for a given subtask. Each subtask is explained in detail in the following sections.
SiameseOI requires the data to be provided as image folders and CSV files. See the Data Annotation Format page for more information about the input data format for SiameseOI.
Configuring a Custom Dataset
This section provides an example configuration and commands for training SiameseOI using the dataset format described above.
You will need to configure the augmentation_config mean and standard deviation based on your input dataset.
Here is an example spec file for training a SiameseOI model with a custom backbone on a custom dataset using the Data Annotation Format.
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
output_shape:
- 100
- 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: 50
checkpoint_interval: 5
results_dir: "${results_dir}/train"
| Parameter | Data Type | Default | Description |
model |
dict config | – | The configuration of the model architecture |
dataset |
dict config | – | The configuration for the dataset detailed in the Config section |
train |
dict config | – | The configuration for training parameters, which is detailed in the Train section |
results_dir |
string | – | The path to save the model experiment log outputs and model checkpoints |
results_dir |
string | – | The path to save the model training experiment log outputs and model checkpoints |
checkpoint_interval |
int | 5 | The interval at which the checkpoint needs to be saved |
num_epochs |
int | 20 | The total number of epochs to run the experiment |
optim |
dict config | None | Contains the configurable parameters for the SiameseOI optimizer detailed in the optim section. |
loss |
str | contrastive | The loss function used during training |
optim
optim:
lr: 0.0005
| Parameter | Datatype | Default | Description | Supported Values |
lr |
float | 0.0005 | The learning rate | >=0.0 |
The following example model config provides options to change the SiameseOI architecture for training.
model:
model_type: Siamese_3
model_backbone: custom
embedding_vectors: 5
margin: 2.0
The following example model is used during SiameseOI evaluation/inference.
| Parameter | Datatype | Default | Description | Supported Values |
model_type |
string | Siamese_3 | The default model architecture from the supported custom model architectures | Siamese_3, Siamese_1 |
model_backbone |
string | custom | The name of the backbone to use | custom |
embedding_vectors |
int | 5 | The embedding dimensions of the final output from the model before computing Euclidian distance | |
margin |
float | 2.0 | The threshold parameter that determines the minimum distance between embeddings of positive and negative pairs |
The dataset parameter defines the dataset source, training batch size,
augmentation, and pre-processing. An example dataset is provided below.
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
output_shape:
- 100
- 100
augmentation_config:
rgb_input_mean: [0.485, 0.456, 0.406]
rgb_input_std: [0.229, 0.224, 0.225]
| Parameter | Datatype | Default | Description | Supported Values |
train_dataset |
Dict | – | The paths to the image directory and CSV files for the training dataset | |
validation_dataset |
Dict | – | The paths to the image directory and CSV files for the validation dataset | |
image_ext |
str | .jpg | The file extension of the images in the dataset | string |
batch_size |
int | 32 | The number of samples per batch | string |
workers |
int | 8 | The number of worker processes for data loading | |
fpratio_sampling |
int | 0.1 | The ratio of false-positive examples to sample | >0 |
num_input |
int | 4 | The number of lighting conditions for each input image* | >0 |
input_map |
Dict | – | The mapping of lighting conditions to indices specifying concatenation ordering* | |
concat_type |
string | linear | Type of concatenation to use for different image lighting conditions | linear, grid |
|
|
Dict Dict dict config |
None None None |
The parameters to define the grid dimensions to concatenate images as a grid: * x: The number of images along the x-axis * y: The number of images along the y-axis |
Dict |
output_shape |
List[int] | [100, 100] | Image resolution of each lighting condition to be reshaped before concatenation | >=0 |
|
|
Dict List[float] List[float] |
None [0.485, 0.456, 0.406] [0.229, 0.224, 0.225] |
The image normalization config, which contains the following parameters:
*
* |
>=0.0 >=0.0 |
* See the Dataset Annotation Format definition for more information about specifying lighting conditions.
Use the following command to run SiameseOI training:
tao model optical_inspection train -e <experiment_spec_file>
-r <results_dir>
--gpus <num_gpus>
Required Arguments
-e, --experiment_spec_file: The path to the experiment spec file-r, --results_dir: The path to a folder where the experiment outputs should be written
Optional Arguments
--gpus: The number of GPUs to use for training. The default value is 1.
Here’s an example of using the SiameseOI training command:
tao model optical_inspection train -e $DEFAULT_SPEC -r $RESULTS_DIR --gpus $NUM_GPUs
Here is an example spec file for testing evaluation and inference of a trained SiameseOI model.
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
output_shape:
- 100
- 100
augmentation_config:
rgb_input_mean: [0.485, 0.456, 0.406]
rgb_input_std: [0.229, 0.224, 0.225]
evaluate:
gpu_id: 0
checkpoint: "${results_dir}/train/oi_model_epoch=019.pth"
results_dir: "${results_dir}/evaluate"
inference:
gpu_id: 0
checkpoint: "${results_dir}/train/oi_model_epoch=004.pth"
results_dir: "${results_dir}/inference"
Use the following command to run SiameseOI evaluation:
tao model optical_inspection evaluate -e <experiment_spec>
-r <results_dir>
Required Arguments
-e, --experiment_spec_file: The experiment spec file to set up the evaluation experiment-r, --results_dir: The path to a folder where the experiment outputs should be written
Here’s an example of using the SiameseOI evaluation command:
tao model optical_inspection evaluate -e $DEFAULT_SPEC -r $RESULTS_DIR
Use the following command to run inference on SiameseOI with the .tlt model:
tao model optical_inspection inference -e <experiment_spec>
-r <results_dir>
Required Arguments
-e, --experiment_spec_file: The experiment spec file to set up the evaluation experiment-r, --results_dir: The path to a folder where the experiment outputs should be written
Here’s an example of using the SiameseOI inference command:
tao model optical_inspection inference -e $DEFAULT_SPEC -r $RESULTS_DIR
Here is an example spec file for exporting the trained SiameseOI model:
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
Use the following command to export the model:
tao model optical_inspection export [-h] -e <experiment spec file>
-r <results_dir>
Required Arguments
-e, --experiment_spec_file: The experiment spec file to set up the evaluation experiment-r, --results_dir: The path to a folder where the experiment outputs should be written
Sample Usage
The following is an example export command:
tao model optical_inspection export -e /path/to/spec.yaml -r $RESULTS_DIR
For deployment, refer to the TAO Deploy Documentation for SiameseOI.