Mask Auto Labeler

Mask Auto Labeler (MAL) is a high-quality, transformer-based mask auto-labeling framework for instance segmentation using only box annotations. It supports the following tasks:

• train

• evaluate

• inference

These tasks may be invoked from the TAO Launcher using the following convention on the command line:

tao mal <sub_task> <args_per_subtask>

Where args_per_subtask are the command-line arguments required for a given subtask. Each of these subtasks are explained in detail below.

Creating a Configuration File

FTMS ClientTAO Launcher

SPECS=$(tao-client mal get-spec --action train --job_type experiment --id $EXPERIMENT_ID)

See also

For information on how to create an experiment using the remote client, refer to the Creating an experiment section in the Remote Client documentation.

Parameter	Datatype	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
strategy	string	‘ddp’	The distributed training strategy	‘ddp’, ‘fsdp’

Dataset Config

The dataset configuration (dataset) defines the data source and input size.

Field	Datatype	Default	Description	Supported Values
train_ann_path	string	–	The path to the training annotation JSON file
val_ann_path	string	–	The path to the validation annotation JSON file
train_img_dir	string	–	The path to the training image directory
val_img_dir	string	–	The path to the validation annotation JSON file
crop_size	Unsigned int	512	The effective input size of the model
load_mask	boolean	True	A flag specifying whether to load the segmentation mask from the JSON file
min_obj_size	float	2048	The minimum object size for training
max_obj_size	float	1e10	The maximum object size for training
num_workers_per_gpu	Unsigned int		The number of workers to load data for each GPU

Model Config

The model configuration (model) defines the model architecture.

Field	Datatype	Default	Description	Supported Values
arch	string	vit-mae-base/16	The backbone architecture Supported backbones include the following: vit-deit-tiny/16 vit-deit-small/16 vit-mae-base/16 vit-mae-large/16 vit-mae-huge/14 fan_tiny_12_p16_224 fan_small_12_p16_224 fan_base_18_p16_224 fan_large_24_p16_224 fan_tiny_8_p4_hybrid fan_small_12_p4_hybrid fan_base_16_p4_hybrid fan_large_16_p4_hybrid
frozen_stages	List[int]	[-1]	The indices of the frozen blocks
mask_head_num_convs	Unsigned int	4	The number of conv layers in the mask head
mask_head_hidden_channel	Unsigned int	256	The number of conv channels in the mask head
mask_head_out_channel	Unsigned int	256	The number of output channels in the mask head
teacher_momentum	float	0.996	The momentum of the teacher model

Train Config

The training configuration (train) specifies the parameters for the training process.

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
batch_size	Unsigned int		The training batch size
use_amp	boolean	True	A flag specifying whether to use mixed precision
optim_momentum	float	0.9	The momentum of the AdamW optimizer
lr	float	0.0000015	The learning rate
min_lr_rate	float	0.2	The minimum learning rate ratio
wd	float	0.0005	The weight decay
warmup_epochs	Unsigned int	1	The number of epochs for warmup
crf_kernel_size	Unsigned int	3	The kernel size of the mean field approximation
crf_num_iter	Unsigned int	100	The number of iterations to run mask refinement
loss_mil_weight	float	4	The weight of multiple instance learning loss
loss_crf_weight	float	0.5	The weight of conditional random field loss

Evaluation Config

The evaluation configuration (evaluate) specifies the parameters for the validation during training as well as the standalone evaluation.

Field	Datatype	Default	Description	Supported Values
checkpoint	string		Path to PyTorch model to evaluate
results_dir	string	/results/evaluate	The directory to save evaluation results
num_gpus	unsigned int	1	The number of GPUs to use for distributed evaluation	>0
gpu_ids	List[int]	[0]	The indices of the GPU’s to use for distributed evaluation
batch_size	Unsigned int		The evaluation batch size
use_mixed_model_test	boolean	False	A flag specifying whether to evaluate with a mixed model
use_teacher_test	boolean	False	A flag specifying whether to evaluate with the teacher model

Inference Config

The inference configuration (inference) specifies the parameters for generating pseudo masks given the groundtruth bounding boxes in COCO format.

Field	Datatype	Default	Description	Supported Values
checkpoint	string		Path to PyTorch model to inference
results_dir	string	/results/inference	The directory to save inference results
num_gpus	unsigned int	1	The number of GPUs to use for distributed inference	>0
gpu_ids	List[int]	[0]	The indices of the GPU’s to use for distributed inference
ann_path	string		The path to the annotation JSON file
img_dir	string		The image directory
label_dump_path	string		The path to save the output JSON file with pseudo masks
batch_size	Unsigned int		The inference batch size
load_mask	boolean	False	A flag specifying whether to load masks if the annotation file has them

Training the Model

Use the following command to run MAL training:

FTMS ClientTAO Launcher

TRAIN_JOB_ID=$(tao-client mal experiment-run-action --action train --id $EXPERIMENT_ID --specs "$SPECS")

See also

For information on how to create an experiment using the remote client, refer to the Creating an experiment section in the Remote Client documentation.

Evaluating the Model

To run evaluation for a MAL model, use this command:

FTMS ClientTAO Launcher

EVAL_JOB_ID=$(tao-client mal experiment-run-action --action evaluate --id $EXPERIMENT_ID --parent_job_id $TRAIN_JOB_ID --specs "$SPECS")

See also

For information on how to create an experiment using the remote client, refer to the Creating an experiment section in the Remote Client documentation.

Running Inference

The inference tool for MAL networks can be used to generate pseudo masks. Here’s an example of using this tool:

FTMS ClientTAO Launcher

INFERENCE_JOB_ID=$(tao-client mal experiment-run-action --action inference --id $EXPERIMENT_ID --parent_job_id $TRAIN_JOB_ID --specs "$SPECS")

See also

For information on how to create an experiment using the remote client, refer to the Creating an experiment section in the Remote Client documentation.