SegFormer
SegFormer 是 NVIDIA 开发的语义分割模型,包含在 TAO 中。SegFormer 支持以下任务
训练
评估
推理
导出
这些任务可以使用 TAO Launcher 从命令行调用,使用以下约定
tao model segformer <sub_task> <args_per_subtask>
其中 args_per_subtask
是给定子任务所需的命令行参数。以下章节详细解释了每个子任务。
Segformer 要求数据以图像和掩码文件夹的形式提供。有关 Segformer 输入数据格式的更多信息,请参阅数据标注格式页面。
自定义数据集的配置
在本文档中,我们展示了在 ISBI 数据集上进行训练的示例配置和命令。ISBI 挑战赛:EM 堆栈神经元结构分割数据集,用于二元分割。它包含灰度图像。有关更多详细信息,请参阅示例 notebook TAO 计算机视觉 示例。因此,我们将 :code: input_type 设置为
grayscale
。对于“RGB”输入图像,:code: input_type 应设置为
rgb
而不是grayscale
。请根据您的输入数据集配置
img_norm_cfg
的均值和标准差。
以下是在 ISBI 数据集上使用 mit_b5 主干训练 SegFormer 模型的示例规范文件。
train:
exp_config:
manual_seed: 49
checkpoint_interval: 200
logging_interval: 50
max_iters: 1000
resume_training_checkpoint_path: null
validate: True
validation_interval: 500
trainer:
find_unused_parameters: True
sf_optim:
lr: 0.00006
model:
input_height: 512
input_width: 512
pretrained_model_path: null
backbone:
type: "mit_b5"
dataset:
data_root: /tlt-pytorch
input_type: "grayscale"
img_norm_cfg:
mean:
- 127.5
- 127.5
- 127.5
std:
- 127.5
- 127.5
- 127.5
to_rgb: True
train_dataset:
img_dir:
- /data/images/train
ann_dir:
- /data/masks/train
pipeline:
augmentation_config:
random_crop:
cat_max_ratio: 0.75
resize:
ratio_range:
- 0.5
- 2.0
random_flip:
prob: 0.5
val_dataset:
img_dir: /data/images/val
ann_dir: /data/masks/val
palette:
- seg_class: foreground
rgb:
- 0
- 0
- 0
label_id: 0
mapping_class: foreground
- seg_class: background
rgb:
- 255
- 255
- 255
label_id: 1
mapping_class: background
repeat_data_times: 500
batch_size: 4
workers_per_gpu: 1
训练分类实验规范由三个主要组件组成
训练
dataset
model
训练配置包含与训练相关的参数。 它们描述如下
参数 | 数据类型 | 默认值 | 描述 | 支持的值 |
exp_config |
Dict int |
None 49 |
* 用于使训练具有确定性的随机种子 |
– |
max_iters |
int | 10 | 训练应进行的最大迭代次数/步数 | |
checkpoint_interval |
int | 1 | 需要保存检查点的步数 | |
logging_interval |
int | 10 | 需要保存实验日志的步数。 日志保存在 logs 目录中。 | |
resume_training_checkpoint_path |
str | None | 用于恢复训练的检查点路径 | |
validate |
bool | False | 在训练期间启用验证的标志 | |
|
int |
int |
在训练期间应执行验证的迭代间隔数 |
|
|
Dict bool |
None False |
此配置包含 MMSeg trainer 所需的参数 * |
– |
sf_optim
sf_optim:
lr: 0.00006
betas:
- 0.0
- 0.999
paramwise_cfg:
pos_block:
decay_mult: 0.0
norm:
decay_mult: 0.0
head:
lr_mut: 10.0
weight_decay: 5e-4
参数 | 数据类型 | 默认值 | 描述 | 支持的值 |
lr |
float | 0.00006 | 学习率 | >=0.0 |
betas |
List[float] | [0.0, 0.9] | Adam 优化器中的 beta 参数 | >=0.0 |
|
Dict |
None |
Adam 优化器的配置参数 * *
|
– |
weight_decay |
float | 5e-4 | 用于正则化的 weight_decay 超参数。 | >=0.0 |
lr_config
lr_config:
warmup_iters: 1500
warmup_ratio: 1e-6
power: 1.0
min_lr: 0.0
参数 | 数据类型 | 默认值 | 描述 | 支持的值 |
warmup_iters |
int | 1500 | 预热持续的迭代或 epoch 数。 | >=0.0 |
warmup_ratio |
float | 1e-6 | 预热开始时使用的 LR 等于 warmup_ratio * initial_lr |
>=0.0 |
power |
float | 1.0 | 将乘法系数提高到的幂。 | >=0.0 |
min_lr |
float | 0.0 | 启动 LR 调度器的最小 LR | >=0.0 |
以下示例 model
提供了更改 SegFormer 架构以进行训练的选项。
model:
input_height: 512
input_width: 512
pretrained_model_path: null
backbone:
type: "mit_b5"
以下示例 model
用于 Segformer 评估/推理期间。
参数 | 数据类型 | 默认值 | 描述 | 支持的值 |
pretrained_model_path |
string | None | 预训练主干文件的可选路径 | string to the path |
|
Dict |
None |
包含以下可配置参数的字典 |
mit_b0, mit_b1 |
|
Dict Float |
None 0.1 |
包含解码器参数的字典 *
|
256, 512, 768 >=0.0 |
input_width |
int | 512 | 模型的输入高度 | >0 |
input_height |
int | 512 | 模型的输入宽度 | >0 |
dataset
参数定义数据集源、训练批次大小和增强。下面提供了 dataset
的示例。
dataset:
data_root: /tlt-pytorch
input_type: "grayscale"
img_norm_cfg:
mean:
- 127.5
- 127.5
- 127.5
std:
- 127.5
- 127.5
- 127.5
to_rgb: True
train_dataset:
img_dir:
- /data/images/train
ann_dir:
- /data/masks/train
pipeline:
augmentation_config:
random_crop:
cat_max_ratio: 0.75
resize:
ratio_range:
- 0.5
- 2.0
random_flip:
prob: 0.5
val_dataset:
img_dir: /data/images/val
ann_dir: /data/masks/val
palette:
- seg_class: foreground
rgb:
- 0
- 0
- 0
label_id: 0
mapping_class: foreground
- seg_class: background
rgb:
- 255
- 255
- 255
label_id: 1
mapping_class: background
repeat_data_times: 500
batch_size: 4
workers_per_gpu: 1
参数 | 数据类型 | 默认值 | 描述 | 支持的值 |
|
Dict |
None |
mage 归一化配置,其中包含以下参数 |
>=0, <=255 |
input_type |
String | “rgb” | 输入类型是 RGB 还是灰度 | “rgb”, “grayscale” |
|
List[Dict] |
None |
调色板配置
|
string |
batch_size |
unsigned int | 32 | 训练和验证的批次大小 | >0 |
workers_per_gpu |
unsigned int | 8 | 并行处理数据的工作进程数 | >0 |
|
dict config str dict config |
None None None |
用于定义训练数据集的参数 * *
|
Dict Config None |
|
dict config str |
None None |
验证配置包含以下用于验证的参数
|
>=0 |
|
dict config str |
None None |
验证配置包含以下用于验证的参数
|
>=0 |
augmentation_config
参数 | 数据类型 | 默认值 | 描述 | 支持的值 |
|
Dict |
None |
random_crop 配置具有以下参数
|
0< h,w <= img_ht, img_wd |
|
Dict Bool |
None [0.5, 2.0] True |
resize Config 具有以下可配置参数 *
|
>=0 True/ False |
|
Dict |
None |
random_flip 配置包含以下用于翻转增强的参数
|
>=0.0 |
使用以下命令运行 Segformer 训练
tao model segformer train [-h] -e <experiment_spec_file>
[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.train_config.optimizer.<optim_option>
: 优化器选项
对于训练、评估和推理,我们为每个相应的任务公开 2 个变量:num_gpus
和 gpu_ids
,默认值分别为 1
和 [0]
。如果两者都已传递,但不一致,例如 num_gpus = 1
、gpu_ids = [0, 1]
,则会修改它们以遵循具有更多 GPU 的设置,例如 num_gpus = 1 -> num_gpus = 2
。
Segformer 的评估指标是 meanIOU。有关 mean IOU 指标的更多详细信息,请参阅此处 meanIOU。
使用以下命令运行 Segformer 评估
tao model segformer evaluate -e <experiment_spec>
evaluate.checkpoint=<evaluation model>
results_dir=<path to output evaluation results>
[evaluate.gpu_ids=<gpu indices>]
[evaluate.num_gpus=<number of gpus>]
必需参数
-e, --experiment_spec_file
: 用于设置评估实验的实验规范文件。evaluate.checkpoint
:.pth
模型。
以下是使用 Segformer 评估命令的示例
+------------+-------+-------+
| Class | IoU | Acc |
+------------+-------+-------+
| foreground | 37.81 | 44.56 |
| background | 83.81 | 95.51 |
+------------+-------+-------+
Summary:
+--------+-------+-------+-------+
| Scope | mIoU | mAcc | aAcc |
+--------+-------+-------+-------+
| global | 60.81 | 70.03 | 85.26 |
+--------+-------+-------+-------+
...
使用以下命令在 Segformer 上使用 .pth
模型运行推理。
tao model segformer inference -e <experiment_spec>
inference.checkpoint=<inference model>
results_dir=<path to output directory for inference>
[inference.gpu_ids=<gpu indices>]
[inference.num_gpus=<number of gpus>]
带有类别 ID 的输出掩码 PNG 图像保存在 vis_tao
中。叠加的掩码图像保存在 mask_tao
中。
必需参数
-e, --experiment_spec
: 用于设置推理的实验规范文件inference.checkpoint
: 用于执行推理的.pth
模型results_dir
: 用于保存推理掩码和掩码叠加图像的路径。推理创建两个目录。
使用以下命令导出模型。
tao model segformer export [-h] -e <experiment spec file>
results_dir=<path to results dir>
export.checkpoint=<trained pth model to be xported>
export.onnx_file=<onnx path>
必需参数
-e, --experiment_spec
: 实验规范文件的路径results_dir
: 将保存导出日志的路径export.checkpoint
: 要导出的.pth
模型export.onnx_file
: 要存储的 :code:.`onnx` 文件
有关部署,请参阅 TAO Deploy 文档
有关将 SegFormer 模型部署到 DeepStream 的更多信息,请参阅集成 SegFormer 模型页面。