mask2former - NVIDIA 文档

Mask2Former 支持以下任务

训练
评估
推理
导出

这些任务可以使用以下命令行约定从 TAO Launcher 调用

复制
已复制！

            
            tao model mask2former <sub_task> <args_per_subtask>

其中 args_per_subtask 是给定子任务所需的命令行参数。以下将解释每个子任务。

数据集格式

Mask2Former 支持 3 种类型的数据加载器，分别对应于语义分割、全景分割和实例分割任务。

每个数据加载器都需要特定的注释格式。

对于语义分割任务，JSONL 注释文件的每一行都编码了原始图像和掩码真值的位置。

对于全景分割和实例分割任务，注释格式分别遵循 COCO 全景格式和 COCO 格式。

注意

类别 ID 和注释 ID 必须大于 0。

创建配置文件

以下是 Mask2Former 规范文件的示例。它有六个组件——model、inference、evaluate、dataset、export 和 train——以及几个全局参数，如下所述。规范文件的格式为 YAML 文件。

这是 Mask2Former 规范文件的示例

复制
已复制！

            
            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`	dict config	–	模型架构的配置
`dataset`	dict config	–	数据集的配置
`训练`	dict config	–	训练任务的配置
`评估`	dict config	–	评估任务的配置
`推理`	dict config	–	推理任务的配置
`encryption_key`	string	无	用于加密和解密模型文件的加密密钥
`results_dir`	string	/results	实验结果保存目录
`导出`	dict config	–	ONNX 导出任务的配置

模型配置

模型配置 (model) 定义了 Mask2Former 模型结构。此模型用于训练、评估和推理。下表包含详细描述。目前，Mask2Former 仅支持 Swin-Transformers 和 EfficientViT（实验性功能）模型。

字段	描述	数据类型和约束	支持的值
backbone	主干网络配置	Dict
sem_seg_head	分割头的配置	Dict
mask_former	mask2former 架构的配置	Dict
mode	后处理模式	string	‘全景’, ‘语义’, ‘实例’
object_mask_threshold	分类置信度阈值	float	0.4
overlap_threshold	全景推理的重叠阈值	float	0.8
test_topk_per_image	为实例推理保留每张图像的 topk 个实例	无符号整数	100

主干网络配置

主干网络配置 (backbone) 定义了主干网络结构。下表包含详细描述。目前，Mask2Former 仅支持 Swin-Transformers 和 EfficientViT 模型。

字段	描述	数据类型和约束	推荐/典型值
type	主干网络类型	str	“swin”
pretrained_weights	预训练主干网络模型的路径	str
swin	Swin 主干网络的配置	Dict
efficientvit	EfficientViT 主干网络的配置	Dict

Swin 配置

swin 配置 (swin) 指定了 Swin Transformer 主干网络中的关键参数。

字段	描述	数据类型和约束	推荐/典型值
type	Swin Transformer 的类型（从 `tiny` 到 `huge`）	str	“large”
pretrain_img_size	预训练中使用的图像大小	无符号整数	384
out_indices	提取特征图的阶段	List	[0, 1, 2, 3]
out_features	提取的特征图的名称	List	[“res2”, “res3”, “res4”, “res5”]

EfficientViT 配置

efficientvit 配置 (efficientvit) 指定了 EfficientViT 主干网络中的关键参数。

字段	描述	数据类型和约束	推荐/典型值
name	EfficientViT 模型的名称（“b0”-“b3”, “l0”-“l3”）	str	“l2”
pretrain_img_size	预训练中使用的图像大小	无符号整数	384
out_indices	提取特征图的阶段	List	[0, 1, 2, 3]
out_features	提取的特征图的名称	List	[“res2”, “res3”, “res4”, “res5”]

数据配置

数据配置 (data) 定义了数据源、增强方法和预处理超参数。

字段	描述	数据类型和约束	推荐/典型值
pixel_mean	RGB 顺序的图像均值	List	[0.485, 0.456, 0.406]
pixel_std	RGB 顺序的图像标准差	List	[0.229, 0.224, 0.225]
augmentation	增强设置	Dict
contiguous_id	是否使用连续 ID	bool
label_map	标签映射文件的路径	string
workers	每个 GPU 加载数据的工作线程数	无符号整数
训练	train	Dict
val	验证数据集配置	Dict
test	测试数据集配置	Dict

增强配置

增强配置 (augmentation) 定义了增强方法。

参数	数据类型	描述	支持的值
`train_min_size`	int list	用于训练数据执行随机调整大小的尺寸列表	int list
`train_max_size`	无符号整数	训练数据的最小随机裁剪尺寸	>0
`train_crop_size`	int list	训练数据的随机裁剪尺寸，格式为 [H, W]	int list
`test_min_size`	无符号整数	测试数据的最小调整大小尺寸	>0
`test_max_size`	无符号整数	测试数据的最大调整大小尺寸	>0

数据集配置

数据集配置 (dataset) 定义了 train、val 或 test 的数据集目录、注释文件和批次大小。

参数	数据类型	描述
`type`	str	数据集类型 (“ade”, “coco”, “coco_panoptic”)
`panoptic_json`	str	COCO 全景格式的 JSON 文件
`img_dir`	str	图像目录（可以是 `root_dir` 的相对路径）
`panoptic_dir`	str	全景分割注释图像目录
`root_dir`	str	`img_dir` 的根目录
`annot_file`	str	COCO/COCO_panoptic 格式的 JSON 文件，或图像/掩码对的 JSONL 格式
`batch_size`	无符号整数	批次大小
`num_workers`	无符号整数	处理输入数据的工作线程数

训练配置

train 配置定义了训练过程的超参数。

复制
已复制！

            
            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

参数	数据类型	默认值	描述	支持的值
`num_gpus`	无符号整数	1	用于分布式训练的 GPU 数量	>0
`gpu_ids`	List[int]	[0]	用于分布式训练的 GPU 索引列表
`seed`	无符号整数	1234	random、NumPy 和 torch 的随机种子	>0
`num_epochs`	无符号整数	10	运行实验的总 epoch 数	>0
`checkpoint_interval`	无符号整数	1	保存检查点的 epoch 间隔	>0
`validation_interval`	无符号整数	1	运行验证的 epoch 间隔	>0
`resume_training_checkpoint_path`	string		从中恢复训练的中间 PyTorch Lightning 检查点
`results_dir`	string	/results/train	保存训练结果的目录
`optim`	dict config		优化器的配置，包括学习率、学习率调度器和权重衰减	>0
`clip_grad_type`	str	full	梯度裁剪方法的类型
`clip_grad_norm`	float	0.1	按 L2 范数裁剪梯度的量。值为 0.0 表示不裁剪	>=0
`precision`	string	fp32	指定 “fp16” 启用混合精度训练。使用 fp16 训练可以帮助节省 GPU 内存。	fp32, fp16
`distributed_strategy`	string	ddp	多 GPU 训练策略。支持 DDP（分布式数据并行）和 Sharded DDP。	ddp, ddp_sharded
`activation_checkpoint`	bool	True	True 值指示 train 在反向传播中重新计算以节省 GPU 内存，而不是存储激活。	True, False
`pretrained_model_path`	string		用于微调的预训练模型检查点路径
`num_nodes`	无符号整数	1	节点数。如果值大于 1，则启用多节点	>0
`freeze`	string list	[]	模型中要冻结的图层名称列表。示例 [“backbone”, “transformer.encoder”, “input_proj”]
`verbose`	bool	False	是否打印来自优化器的详细学习率缩放	True, False
`iters_per_epoch`	无符号整数		每个 epoch 的样本数

优化器配置

optim 参数定义了训练中优化器的配置，包括学习率、学习率调度器和权重衰减。

参数	数据类型	默认值	描述	支持的值
`lr`	float	2e-4	训练模型（不包括主干网络）的初始学习率	>0.0
`momentum`	float	0.9	AdamW 优化器的动量	>0.0
`weight_decay`	float	1e-4	权重衰减系数	>0.0
`lr_scheduler`	string	MultiStep	学习率调度器 * `MultiStep` : 从 `lr_steps` 开始，按 `lr_decay` 减小 `lr` * `StepLR` : 在每个 `lr_step_size` 时，按 `lr_decay` 减小 `lr`	MultiStep/StepLR
`gamma`	float	0.1	学习率调度器的递减因子	>0.0
`milestones`	int list	[11]	对于 `MultiStep` 调度器，减小学习率的步数	int list
`monitor_name`	string	val_loss	`AutoReduce` 调度器的监控值	val_loss/train_loss
`type`	string	AdamW	训练期间要使用的优化器类型	AdamW/SGD

评估配置

evaluate 参数定义了评估过程的超参数。

复制
已复制！

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

参数	数据类型	默认值	描述	支持的值
`checkpoint`	string		要评估的 PyTorch 模型路径
`trt_engine`	string		要评估的 TensorRT 模型路径。必须仅与 tao deploy 一起使用
`num_gpus`	无符号整数	1	num_gpus	>0
`gpu_ids`	无符号整数	[0]	gpu_ids
`results_dir`	string	/results/evaluate	评估结果目录的路径

推理配置

inference 参数定义了推理过程的超参数。

复制
已复制！

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

参数	数据类型	默认值	描述	支持的值
`checkpoint`	string		checkpoint
`trt_engine`	string		trt_engine
`num_gpus`	无符号整数	1	num_gpus	>0
`gpu_ids`	无符号整数	[0]	gpu_ids
`results_dir`	string	/results/inference	推理结果目录的路径

导出配置

export 参数定义了导出过程的超参数。

复制
已复制！

            
            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

参数	数据类型	默认值	描述	支持的值
`checkpoint`	string		checkpoint_path
`onnx_file`	string		`.onnx` 文件的路径
`on_cpu`	bool	True	如果此值为 True，则 DMHA 模块将导出为标准 PyTorch。如果此值为 False，则模块将使用 TRT 插件导出。	True, False
`opset_version`	无符号整数	12	导出的 ONNX 的 opset 版本	>0
`input_channel`	无符号整数	3	输入通道大小。仅支持值 3。	3
`input_width`	无符号整数	960	输入宽度	>0
`input_height`	无符号整数	544	输入高度	>0
`batch_size`	无符号整数	-1	dynamic_batch_size	>=-1

训练模型

要训练 Mask2Former 模型，请使用此命令

复制
已复制！

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

必需参数

-e, --experiment_spec：用于设置训练实验的实验规范文件。

可选参数

您可以设置可选参数来覆盖实验规范文件中的选项值。

-h, --help：显示此帮助消息并退出。
model.<model_option>：模型选项。
dataset.<dataset_option>：数据集选项。
train.<train_option>：训练选项。
train.optim.<optim_option>：优化器选项

注意

对于训练、评估和推理，我们为每个相应的任务公开 2 个变量：num_gpus 和 gpu_ids，它们分别默认为 1 和 [0]。如果两者都传递，但彼此不一致，例如 num_gpus = 1，gpu_ids = [0, 1]，则它们将被修改以遵循具有更多 GPU 的设置，例如 num_gpus = 1 -> num_gpus = 2。

检查点和恢复训练

在每个 train.checkpoint_interval 时，将保存一个 PyTorch Lightning 检查点。它被称为 model_epoch_<epoch_num>.pth。这些保存在 train.results_dir 中，如下所示

复制
已复制！

            
            $ ls /results/train

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

最新的检查点也将保存为 mask2former_model_latest.pth。如果 mask2former_model_latest.pth 存在于 train.results_dir 中，则训练将自动从它恢复。如果提供了 train.resume_training_checkpoint_path，则它将取代此项。

此逻辑的主要含义是，如果您希望从头开始触发全新训练，则可以

优化 Mask2Former 训练的资源

训练 Mask2Former 需要强大的 GPU（例如，V100/A100），至少具有 15GB 的 VRAM 和大量 CPU 内存，才能在像 COCO 这样的标准数据集上进行训练。本节概述了您可以使用的一些策略，以便仅使用有限的资源启动训练。

优化 GPU 内存

有多种方法可以优化 GPU 内存使用量。一种典型的选择是减少 dataset.batch_size。但是，这可能会导致您的训练时间比平时更长。我们建议设置以下配置以优化 GPU 消耗

将 train.precision 设置为 fp16 以启用自动混合精度训练。这可以减少 50% 的 GPU 内存使用量。
将 train.activation_checkpoint 设置为 True 以启用激活检查点。通过重新计算激活而不是将它们缓存到内存中，可以提高内存使用率。
将 train.distributed_strategy 设置为 ddp_sharded 以启用 Sharded DDP 训练。这会在不同进程之间共享梯度计算，以帮助减少 GPU 内存。
尝试使用更轻量级的主干网络，或通过设置 train.freeze 来冻结主干网络。
尝试根据您的数据集在 dataset.augmentation 中更改增强分辨率。

优化 CPU 内存

为了加快数据加载速度，通常的做法是设置大量工作线程以生成多个进程。但是，如果您的注释文件非常大，这可能会导致您的 CPU 内存不足。我们建议设置以下配置以优化 CPU 消耗。

评估模型

要使用 Mask2Former 模型运行评估，请使用此命令

复制
已复制！

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

必需参数

-e, --experiment_spec：用于设置评估实验的实验规范文件。
evaluate.checkpoint：要评估的 .pth 模型。

可选参数

evaluate.<evaluate_option>：评估选项。

使用 Mask2Former 模型运行推理

Mask2Former 模型的推理工具可用于可视化边界框和掩码。

复制
已复制！

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

必需参数

-e, --experiment_spec：用于设置推理实验的实验规范文件。
inference.checkpoint：要在其上运行推理的 .pth 模型。

可选参数

inference.<inference_option>：推理选项。

导出模型

复制
已复制！

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

必需参数

-e, --experiment_spec：实验规范文件的路径
export.checkpoint：要导出的 .pth 模型。
export.onnx_file：保存 .etlt 或 .onnx 模型的路径。

可选参数

export.<export_option>：导出选项。

TensorRT 引擎生成和验证

有关部署，请参阅 Mask2Former 的 TAO Deploy 文档。

部署到 DeepStream

有关将 Mask2Former 模型部署到 DeepStream 的更多信息，请参阅集成 Mask2Former 模型页面。