mask2former

Mask2Former supports the following tasks:

• train

• evaluate

• inference

• export

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

            

            
tao model mask2former <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 as follows.

Mask2Former supports 3 type of dataloaders corresponding to the semantic, panoptic and instance segmentation tasks.

Each dataloader requires a certain annotation format.

For the semantic segmentation task, each line of the JSONL annotation file encodes the locations of the raw image and the mask groundtruth.

For the panoptic and instance segmentation tasks, the annotation format follows the COCO panoptic and COCO format respectively.

Below is a sample Mask2Former spec file. It has six components –model, inference, evaluate, dataset, export, and train–as well as several global parameters, which are described below. The format of the spec file is a YAML file.

Here’s a sample of the Mask2Former spec file:

            

            
results_dir: /workspace/mask2former_coco_swint
data:
  contiguous_id: False
  label_map: /tlt3_experiments/mask2former_coco_effvit_b2/colormap.json
  type: 'coco_panoptic'
  train:
    panoptic_json: "/datasets/coco/annotations/panoptic_train2017.json"
    img_dir: "/datasets/coco/train2017"
    panoptic_dir: "/datasets/coco/panoptic_train2017"
    batch_size: 16
    num_workers: 20
  val:
    panoptic_json: "/datasets/coco/annotations/panoptic_val2017.json"
    img_dir: "/datasets/coco/val2017"
    panoptic_dir: "/datasets/coco/panoptic_val2017"
    batch_size: 1
    num_workers: 2
    target_size: [1024, 1024]
  test:
    img_dir: /workspace/test_images/
    batch_size: 1
  augmentation:
    train_min_size: [1024]
    train_max_size: 2560
    train_crop_size: [1024, 1024]
    test_min_size: 1024
    test_max_size: 2560
train:
  precision: 'fp16'
  num_gpus: 1
  checkpoint_interval: 1
  validation_interval: 5
  num_epochs: 50
  optim:
    lr_scheduler: "MultiStep"
    milestones: [44, 48]
    type: "AdamW"
    lr: 0.0001
    weight_decay: 0.05
model:
  object_mask_threshold: 0.
  overlap_threshold: 0.8
  mode: "semantic"
  backbone:
    pretrained_weights: "/workspace/mask2former_coco_swint/swin_tiny_patch4_window7_224_22k.pth"
    type: "swin"
    swin:
      type: "tiny"
      window_size: 7
      ape: False
      pretrain_img_size: 224
  mask_former:
    num_object_queries: 100
  sem_seg_head:
    norm: "GN"
    num_classes: 200
inference:
  checkpoint: "/workspace/mask2former_coco_swint/train/model_epoch=049.pth"
evaluate:
  checkpoint: "/workspace/mask2former_coco_swint/train/model_epoch=049.pth"
export:
  checkpoint: "/workspace/mask2former_coco_swint/train/model_epoch=049.pth"
  input_channel: 3
  input_width: 1024
  input_height: 1024
  opset_version: 17
        

Model Config

The model configuration (model) defines the Mask2Former model structure. This model is used for training, evaluation, and inference. A detailed description is included in the table below. Currently, Mask2Former only supports Swin-Transformers and EfficientViT (experimental feature) models.

Backbone Config

The backbone configuration (backbone) defines the backbone structure. A detailed description is included in the table below. Currently, Mask2Former only supports Swin-Transformers and EfficientViT models.

Swin Config

The swin configuration (swin) specifies the key parameters in a Swin Transformer backbone.

EfficientViT Config

The efficientvit configuration (efficientvit) specifies the key parameters in a EfficientViT backbone.

Data Config

The data configuration (data) defines the data source, augmentation methods and pre-processing hyperparameters.

Augmentation Config

The augmentation configuration (augmentation) defines the augmentation methods.

Dataset Config

The dataset configuration (dataset) defines the dataset directories, annotation file and batch size for either train, val or test.

Train Config

The train configuration defines the hyperparameters of the training process.

            

            
train:
  precision: 'fp16'
  num_gpus: 1
  checkpoint_interval: 10
  validation_interval: 10
  num_epochs: 50
  optim:
    type: "AdamW"
    lr: 0.0001
    weight_decay: 0.05
        

Optimizer Config

The optim parameter defines the config for the optimizer in training, including the learning rate, learning scheduler, and weight decay.

Evaluation Config

The evaluate parameter defines the hyperparameters of the evaluation process.

            

            
evaluate:
  checkpoint: /path/to/model.pth
  num_gpus: 1
        

Inference Config

The inference parameter defines the hyperparameters of the inference process.

            

            
inference:
  checkpoint: /path/to/model.pth
  num_gpus: 1
        

Export Config

The export parameter defines the hyperparameters of the export process.

            

            
export:
  checkpoint: /path/to/model.pth
  onnx_file: /path/to/model.onnx
  on_cpu: False
  opset_version: 12
  input_channel: 3
  input_width: 960
  input_height: 544
  batch_size: -1
        

To train a Mask2Former model, use this command:

            

            
tao model mask2former train [-h] -e <experiment_spec>
                      [results_dir=<global_results_dir>]
                      [model.<model_option>=<model_option_value>]
                      [dataset.<dataset_option>=<dataset_option_value>]
                      [train.<train_option>=<train_option_value>]
                      [train.gpu_ids=<gpu indices>]
                      [train.num_gpus=<number of gpus>]
        

Required Arguments

• -e, --experiment_spec: The experiment specification file to set up the training experiment.

Optional Arguments

You can set optional arguments to override the option values in the experiment spec file.

• -h, --help: Show this help message and exit.

• model.<model_option>: The model options.

• dataset.<dataset_option>: The dataset options.

• train.<train_option>: The train options.

• train.optim.<optim_option>: The optimizer options

Checkpointing and Resuming Training

At every train.checkpoint_interval, a PyTorch Lightning checkpoint is saved. It is called model_epoch_<epoch_num>.pth. These are saved in train.results_dir, like so:

            

            
$ ls /results/train

'model_epoch_000.pth'
'model_epoch_001.pth'
'model_epoch_002.pth'
'model_epoch_003.pth'
'model_epoch_004.pth'
        

The latest checkpoint will also be saved as mask2former_model_latest.pth. Training automatically resumes from mask2former_model_latest.pth, if it exists in train.results_dir. This is superseded by train.resume_training_checkpoint_path, if it is provided.

The major implication of this logic is that, if you wish to trigger fresh training from scratch, either:

• Specify a new, empty results directory (Recommended)

• Remove the latest checkpoint from the results directory

Optimizing Resource for Training Mask2Former

Training Mask2Former requires strong GPUs (for example, V100/A100) with at least 15GB of VRAM and a lot of CPU memory to be trained on a standard dataset like COCO. This section outlines some of the strategies you can use to launch training with only limited resources.

Optimize GPU Memory

There are various ways to optimize GPU memory usage. A typical option is to reduce dataset.batch_size. However, this can cause your training to take longer than usual. We recommend setting the following configurations to optimize GPU consumption:

• Set train.precision to fp16 to enable automatic mixed precision training. This can reduce your GPU memory usage by 50%.

• Set train.activation_checkpoint to True to enable activation checkpointing. By recomputing the activations instead of caching them into memory, the memory usage can be improved.

• Set train.distributed_strategy to ddp_sharded to enabled Sharded DDP training. This shares gradient calculation across different processes to help reduce GPU memory.

• Try using more lightweight backbones or freeze the backbone through setting train.freeze.

• Try changing the augmentation resolution in dataset.augmentation depending on your dataset.

Optimize CPU Memory

To speed up data loading, it is a common practice to set high number of workers to spawn multiple processes. However, this can cause your CPU memory to become Out of Memory, if the size of your annotation file is very large. We recommend setting the following configurations to optimize CPU consumption.

To run evaluation with a Mask2Former model, use this command:

            

            
tao model mask2former evaluate [-h] -e <experiment_spec>
                      evaluate.checkpoint=<model to be evaluated>
                      [evaluate.<evaluate_option>=<evaluate_option_value>]
                      [evaluate.gpu_ids=<gpu indices>]
                      [evaluate.num_gpus=<number of gpus>]
        

Required Arguments

• -e, --experiment_spec: The experiment spec file to set up the evaluation experiment.

• evaluate.checkpoint: The .pth model to be evaluated.

Optional Arguments

• evaluate.<evaluate_option>: The evaluate options.

The inference tool for Mask2Former models can be used to visualize bboxes and masks.

            

            
tao model mask2former inference [-h] -e <experiment spec file>
                      inference.checkpoint=<inference model>
                      [inference.<evaluate_option>=<evaluate_option_value>]
                      [inference.gpu_ids=<gpu indices>]
                      [inference.num_gpus=<number of gpus>]
        

Required Arguments

• -e, --experiment_spec: The experiment spec file to set up the inference experiment.

• inference.checkpoint: The .pth model to run inference on.

Optional Arguments

• inference.<inference_option>: The inference options.

            

            
tao model mask2former export [-h] -e <experiment spec file>
                      [results_dir=<results_dir>]
                      export.checkpoint=<model to export>
                      export.onnx_file=<onnx path>
        

Required Arguments

• -e, --experiment_spec: The path to an experiment spec file

• export.checkpoint: The .pth model to export.

• export.onnx_file: The path where the .etlt or .onnx model is saved.

Optional Arguments

• export.<export_option>: The export options.

For deployment, refer to TAO Deploy documentation for Mask2Former.

Refer to the Integrating a Mask2Former Model page for more information about deploying a Mask2Former model to DeepStream.