Mask Auto Labeler - NVIDIA 文档

Mask Auto Labeler (MAL) 是一款高质量、基于 Transformer 的掩码自动标注框架，仅使用框标注即可进行实例分割。它支持以下任务

train
evaluate
inference

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

Copy
已复制!

            
            tao mal <sub_task> <args_per_subtask>

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

创建配置文件

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

Copy
已复制!

            
            strategy: 'fsdp'
results_dir: '/path/to/result/dir'
dataset:
  train_ann_path: '/datasets/coco/annotations/instances_train2017.json'
  train_img_dir: '/datasets/coco/raw-data/train2017'
  val_ann_path: '/coco/annotations/instances_val2017.json'
  val_img_dir: '/datasets/coco/raw-data/val2017'
  load_mask: True
  crop_size: 512
inference:
  ann_path: '/dataset/sample.json'
  img_dir: '/dataset/sample_dir'
  label_dump_path: '/dataset/sample_output.json'
model:
  arch: 'vit-mae-base/16'
train:
  num_epochs: 10
  checkpoint_interval: 5
  validation_interval: 5
  batch_size: 4
  seed: 1234
  num_gpus: 1
  gpu_ids: [0]
  use_amp: True
  optim_momentum: 0.9
  lr: 0.0000015
  min_lr_rate: 0.2
  wd: 0.0005
  warmup_epochs: 1
  crf_kernel_size: 3
  crf_num_iter: 100
  loss_mil_weight: 4
  loss_crf_weight: 0.5

参数	数据类型	默认值	描述	支持的值
`model`	dict config	–	模型架构的配置
`dataset`	dict config	–	数据集的配置
`train`	dict config	–	训练任务的配置
`evaluate`	dict config	–	评估任务的配置
`inference`	dict config	–	推理任务的配置
`encryption_key`	string	None	用于加密和解密模型文件的加密密钥
`results_dir`	string	/results	实验结果保存目录
`strategy`	string	‘ddp’	分布式训练策略	‘ddp’, ‘fsdp’

数据集配置

数据集配置 (dataset) 定义数据源和输入大小。

字段	数据类型	默认值	描述	支持的值
`train_ann_path`	string	–	训练标注 JSON 文件的路径
`val_ann_path`	string	–	验证标注 JSON 文件的路径
`train_img_dir`	string	–	训练图像目录的路径
`val_img_dir`	string	–	验证标注 JSON 文件的路径
`crop_size`	Unsigned int	512	模型的有效输入大小
`load_mask`	boolean	True	一个标志，指定是否从 JSON 文件加载分割掩码
`min_obj_size`	float	2048	训练的最小对象大小
`max_obj_size`	float	1e10	训练的最大对象大小
`num_workers_per_gpu`	Unsigned int		每个 GPU 加载数据的工作线程数

模型配置

模型配置 (model) 定义模型架构。

字段	数据类型	默认值	描述	支持的值
`arch`	string	vit-mae-base/16	骨干架构。支持的骨干网络包括以下 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]	冻结块的索引
`mask_head_num_convs`	Unsigned int	4	掩码头中的 conv 层数
`mask_head_hidden_channel`	Unsigned int	256	掩码头中的 conv 通道数
`mask_head_out_channel`	Unsigned int	256	掩码头中的输出通道数
`teacher_momentum`	float	0.996	教师模型的动量

训练配置

训练配置 (train) 指定训练过程的参数。

参数	数据类型	默认值	描述	支持的值
`num_gpus`	unsigned int	1	用于分布式训练的 GPU 数量	>0
`gpu_ids`	List[int]	[0]	用于分布式训练的 GPU 索引
`seed`	unsigned int	1234	random、numpy 和 torch 的随机种子	>0
`num_epochs`	unsigned int	10	运行实验的总 epoch 数	>0
`checkpoint_interval`	unsigned int	1	保存检查点的 epoch 间隔	>0
`validation_interval`	unsigned int	1	运行验证的 epoch 间隔	>0
`resume_training_checkpoint_path`	string		从中恢复训练的中间 PyTorch Lightning 检查点
`results_dir`	string	/results/train	用于保存训练结果的目录
`batch_size`	Unsigned int		训练批次大小
`use_amp`	boolean	True	一个标志，指定是否使用混合精度
`optim_momentum`	float	0.9	AdamW 优化器的动量
`lr`	float	0.0000015	学习率
`min_lr_rate`	float	0.2	最小学习率比率
`wd`	float	0.0005	权重衰减
`warmup_epochs`	Unsigned int	1	预热的 epoch 数
`crf_kernel_size`	Unsigned int	3	均值场近似的内核大小
`crf_num_iter`	Unsigned int	100	运行掩码细化的迭代次数
`loss_mil_weight`	float	4	多实例学习损失的权重
`loss_crf_weight`	float	0.5	条件随机场损失的权重

评估配置

评估配置 (evaluate) 指定训练期间验证以及独立评估的参数。

字段	数据类型	默认值	描述	支持的值
`checkpoint`	string		要评估的 PyTorch 模型路径
`results_dir`	string	/results/evaluate	用于保存评估结果的目录
`num_gpus`	unsigned int	1	用于分布式评估的 GPU 数量	>0
`gpu_ids`	List[int]	[0]	用于分布式评估的 GPU 索引
`batch_size`	Unsigned int		评估批次大小
`use_mixed_model_test`	boolean	False	一个标志，指定是否使用混合模型进行评估
`use_teacher_test`	boolean	False	一个标志，指定是否使用教师模型进行评估

推理配置

推理配置 (inference) 指定用于生成伪掩码的参数，给定 COCO 格式中的 groundtruth 边界框。

字段	数据类型	默认值	描述	支持的值
`checkpoint`	string		要推理的 PyTorch 模型路径
`results_dir`	string	/results/inference	用于保存推理结果的目录
`num_gpus`	unsigned int	1	用于分布式推理的 GPU 数量	>0
`gpu_ids`	List[int]	[0]	用于分布式推理的 GPU 索引
`ann_path`	string		标注 JSON 文件的路径
`img_dir`	string		图像目录
`label_dump_path`	string		用于保存带有伪掩码的输出 JSON 文件的路径
`batch_size`	Unsigned int		推理批次大小
`load_mask`	boolean	False	一个标志，指定如果标注文件有掩码，是否加载它们

训练模型

使用以下命令运行 MAL 训练

Copy
已复制!

            
            tao model mal 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>: 训练选项。

注意

对于训练、评估和推理，我们为每个相应的任务公开 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 中，如下所示

Copy
已复制!

            
            $ ls /results/train

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

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

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

评估模型

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

Copy
已复制!

            
            tao model mal evaluate [-h] -e <experiment_spec_file>
              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>: 评估选项。

运行推理

MAL 网络的 inference 工具可用于生成伪掩码。以下是使用此工具的示例

Copy
已复制!

            
            tao model mal inference [-h] -e <experiment spec file>
              inference.checkpoint=<model to be inferenced>
              [inference.<inference_option>=<inference_option_value>]
              [inference.gpu_ids=<gpu indices>]
              [inference.num_gpus=<number of gpus>]

必需参数

-e, --experiment_spec: 用于设置推理实验的实验规格文件。
inference.checkpoint: 要推理的 .pth 模型。

可选参数

inference.<inference_option>: 推理选项。