SiameseOI#

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

train
evaluate
inference
export

Each task is explained in detail in the following sections.

Data Input for SiameseOI#

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.

Creating a Training Experiment Specification File#

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 specification file for training a SiameseOI model with a custom backbone on a custom dataset using the Data Annotation Format.

Parameter	Data Type	Default	Description	Supported Values
`model`	dict config	–	The configuration of the model architecture
`dataset`	dict config	–	The configuration of the dataset
`train`	dict config	–	The configuration of the training task
`evaluate`	dict config	–	The configuration of the evaluation task
`inference`	dict config	–	The configuration of the inference task
`encryption_key`	string	None	The encryption key to encrypt and decrypt model files
`results_dir`	string	/results	The directory where experiment results are saved
`export`	dict config	–	The configuration of the ONNX export task
`gen_trt_engine`	dict config	–	The configuration of the TensorRT generation task

train#

Parameter	Datatype	Default	Description	Supported Values
`num_gpus`	unsigned int	1	The number of GPUs to use for distributed training	>0
`gpu_ids`	List[int]	[0]	The indices of the GPU’s to use for distributed training
`seed`	unsigned int	1234	The random seed for random, NumPy, and torch	>0
`num_epochs`	unsigned int	10	The total number of epochs to run the experiment	>0
`checkpoint_interval`	unsigned int	1	The epoch interval at which the checkpoints are saved	>0
`validation_interval`	unsigned int	1	The epoch interval at which the validation is run	>0
`resume_training_checkpoint_path`	string		The intermediate PyTorch Lightning checkpoint to resume training from
`results_dir`	string	/results/train	The directory to save training results
`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

Model#

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

Dataset#

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
  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]

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
`input_width`	int	100	The width of the input image	>0
`input_height`	int	100	The height of the input image	>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.

Evaluating the Model#

Multi-GPU evaluation is currently not supported for Optical Inspection.