NVIDIA TAO Toolkit v4.0.1
NVIDIA TAO Release 4.0.1

SiameseOI

SiameseOI is an NVIDIA-developed optical inspection model for PCB data and is included in the TAO Toolkit. SiameseOI supports the following tasks:

  • train

  • evaluate

  • inference

  • export

These tasks can be invoked from the TAO Toolkit Launcher using the following convention on the command-line:

Copy
Copied!
            

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.

Copy
Copied!
            

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

Copy
Copied!
            

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.

Copy
Copied!
            

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.

Copy
Copied!
            

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

grid_map

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

augmentation_config

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:

* rgb_input_mean: The mean to be subtracted for pre-processing

* rgb_input_std: The standard deviation to divide the image by

>=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:

Copy
Copied!
            

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:

Copy
Copied!
            

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.

Copy
Copied!
            

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:

Copy
Copied!
            

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:

Copy
Copied!
            

tao model optical_inspection evaluate -e $DEFAULT_SPEC -r $RESULTS_DIR


Use the following command to run inference on SiameseOI with the .tlt model:

Copy
Copied!
            

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:

Copy
Copied!
            

tao model optical_inspection inference -e $DEFAULT_SPEC -r $RESULTS_DIR


Here is an example spec file for exporting the trained SiameseOI model:

Copy
Copied!
            

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:

Copy
Copied!
            

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:

Copy
Copied!
            

tao model optical_inspection export -e /path/to/spec.yaml -r $RESULTS_DIR


© Copyright 2023, NVIDIA.. Last updated on Jul 27, 2023.