图像分类 PyT
图像分类 PyT 是 TAO 中包含的基于 PyTorch 的图像分类模型。它支持以下任务
训练
评估
推理
导出
这些任务可以使用命令行上的以下约定从 TAO Launcher 调用
tao model classification_pyt <sub_task> <args_per_subtask>
其中，args_per_subtask 是给定子任务所需的命令行参数。以下各节详细解释了每个子任务。
TAO 已更新到最新的 MMPretrain 版本，而不是已弃用的 MMClassification。 .. Note:: 图像分类 (PyT) 基于 MMPretrain。因此，大多数参数都采用 MMPretrain 1.x 格式。
以下是使用 FAN 主干网络的分类 PyT 的数据集规范文件示例
dataset:
data:
samples_per_gpu: 128
workers_per_gpu: 8
train:
data_prefix: "/raid/ImageNet2012/ImageNet2012/train"
pipeline: # Augmentations alone
- type: RandomResizedCrop
scale: 224
- type: RandomFlip
prob: 0.5
direction: "horizontal"
- type: ColorJitter
brightness: 0.4
contrast: 0.4
saturation: 0.4
- type: RandomErasing
erase_prob: 0.3
val:
data_prefix: /raid/ImageNet2012/ImageNet2012/val
test:
data_prefix: /raid/ImageNet2012/ImageNet2012/val
下表描述了 dataset 中的可配置参数。
| 参数 | 数据类型 | 默认值 | 描述 | 支持的值 |
采样器 |
dict 配置 | 无 | 数据集采样器类型 | |
|
|
dict 配置 |
无 |
包含以下可配置参数 |
> 0 |
数据 |
dict 配置 | 无 | 与训练相关的参数。 有关更多详细信息，请参阅 data。 |
数据
| 参数 | 数据类型 | 默认值 | 描述 | 支持的值 |
samples_per_gpu |
整数 | 无 | 数据集采样器类型 | |
|
|
字符串 |
字典 |
包含以下可配置参数 |
> 0 |
|
|
dict 配置 字符串 |
无 |
包含训练数据集配置 |
Imagenet 类 |
|
|
dict 配置 字符串 |
无 |
包含测试数据集配置 |
Imagenet 类 |
|
|
dict 配置 字符串 |
无 |
包含验证数据集配置 |
Imagenet 类 |
有关更多详细信息，请参阅 MMPretrain 1.x 格式文档。
流程
以下是包含不同增强功能的流程配置示例
pipeline: # Augmentations alone
- type: RandomResizedCrop
scale: 224
backend: "pillow"
- type: RandomFlip
prob: 0.5
direction: "horizontal"
- type: ColorJitter
brightness: 0.4
contrast: 0.4
saturation: 0.4
- type: RandomErasing
erase_prob: 0.3
下面列出了一些广泛采用的增强功能和参数。 更多信息，请参阅 MMPretrain 文档，了解转换
| 参数 | 数据类型 | 默认值 | 描述 | 支持的值 |
|
|
dict 配置 |
无 |
包含以下可配置参数 |
> 0 |
|
|
dict 配置 |
无 |
包含以下可配置参数 |
0-1 |
|
|
dict 配置 整数 |
无 |
包含以下可配置参数 |
> 0 |
|
|
dict 配置 |
无 |
ColorJitter 增强功能包含以下参数 |
0-1 |
|
|
dict 配置 |
无 |
RandomErasing 增强功能包含以下参数 |
0-1 |
以下是图像分类 PyT 的训练规范文件示例
train:
train_config:
runner:
max_epochs: 300
checkpoint_config:
interval: 1
logging:
interval: 5000
validate: True
evaluation:
interval: 10
custom_hooks:
- type: "EMAHook"
momentum: 0.00008
priority: "ABOVE_NORMAL"
lr_config:
policy: CosineAnnealing
T_max: 95
by_epoch: True
begin: 5
end: 100
optimizer:
type: AdamW
lr: 0.005
weight_decay: 0.05
下表描述了 train 规范中的可配置参数。
| 参数 | 数据类型 | 默认值 | 描述 | 支持的值 |
|
|
dict 配置 |
无 |
包含以下可配置参数 |
> 0 |
train_config |
dict 配置 | 无 | 与训练相关的参数。 有关更多信息，请参阅 train_config | |
results_dir |
字符串 | 无 | 用于保存检查点和日志的路径 | 字符串 |
train_config
| 参数 | 数据类型 | 默认值 | 描述 | 支持的值 |
|
|
dict 配置 |
无 |
包含以下可配置参数 |
|
|
|
dict 配置 |
无 |
包含以下可配置参数
|
>0 |
|
|
dict 配置 |
无 |
包含以下可配置参数
|
>0 |
|
|
dict 配置 |
无 |
包含不同优化器的可配置参数，详见 |
|
|
|
dict 配置 |
无 |
包含以下参数 |
>=0.0 |
|
|
dict 配置 |
无 |
包含以下可配置参数 |
|
validate |
布尔值 | False | 一个标志，用于在训练期间启用验证 | |
find_unused_parameters |
布尔值 | False | 在 DDP 中设置此参数。 有关更多信息，请参阅 DDP_PyT。 | |
lr_config |
字典 | 无 | 学习率调度器配置。 有关更多详细信息，请参阅 lr_config | |
load_from |
字符串 | 无 | 可以从中加载包括 head 在内的端到端模型权重的检查点路径 | |
custom_hooks |
字典 | 无 | 自定义训练 hooks 配置。 有关更多详细信息，请参阅 custom_hooks。 | |
resume_training_checkpoint_path |
字符串 | 无 | 从其恢复训练的检查点路径 |
optimizer
支持以下优化器
SGD
| 参数 | 数据类型 | 默认值 | 描述 | 支持的值 |
|
|
dict 配置 |
无 |
包含以下可配置参数 |
|
AdamW
| 参数 | 数据类型 | 默认值 | 描述 | 支持的值 |
|
|
dict 配置 |
无 |
包含以下可配置参数 |
|
lr_config
lr_config 参数定义学习率调度器的参数。 下面列出了一些广泛采用的学习率调度器和参数。 更多信息，请参阅 MMPretrain 文档，了解 lr 调度器
CosineAnnealingLR
| 参数 | 数据类型 | 默认值 | 描述 | 支持的值 |
T_max |
整数 | 最大迭代次数。 | >=0.0 | |
by_epoch |
布尔值 | True | 计划的学习率是否按 epoch 更新 | 小于 1.0 |
begin |
整数 | 0 | 开始更新学习率的步数。 | 在区间 (0, INF) 内。 |
end |
整数 | INF | 停止更新学习率的步数。 | 在区间 (0, INF) 内。 |
MultiStepLR
| 参数 | 数据类型 | 典型值 | 描述 | 支持的值 |
gamma |
浮点数 | – | 基本（最大）学习率 | 通常小于 1.0 |
by_epoch |
布尔值 | True | 计划的学习率是否按 epoch 更新 | 小于 1.0 |
begin |
整数 | 0 | 开始更新学习率的步数。 | 在区间 (0, INF) 内。 |
end |
整数 | INF | 停止更新学习率的步数。 | 在区间 (0, INF) 内。 |
LinearLR
| 参数 | 数据类型 | 典型值 | 描述 | 支持的值 |
by_epoch |
布尔值 | True | 计划的学习率是否按 epoch 更新 | 小于 1.0 |
begin |
整数 | 0 | 开始更新学习率的步数。 | 在区间 (0, INF) 内。 |
end |
整数 | INF | 停止更新学习率的步数。 | 在区间 (0, INF) 内。 |
PolyLR
| 参数 | 数据类型 | 典型值 | 描述 | 支持的值 |
eta_min |
浮点数 | 0 | 调度结束时的最小学习率 | |
by_epoch |
布尔值 | True | 计划的学习率是否按 epoch 更新 | 小于 1.0 |
begin |
整数 | 0 | 开始更新学习率的步数。 | 在区间 (0, INF) 内。 |
end |
整数 | INF | 停止更新学习率的步数。 | 在区间 (0, INF) 内。 |
StepLR
| 参数 | 数据类型 | 典型值 | 描述 | 支持的值 |
gamma |
浮点数 | – | 基本（最大）学习率 | 通常小于 1.0 |
by_epoch |
布尔值 | True | 计划的学习率是否按 epoch 更新 | 小于 1.0 |
begin |
整数 | 0 | 开始更新学习率的步数。 | 在区间 (0, INF) 内。 |
end |
整数 | INF | 停止更新学习率的步数。 | 在区间 (0, INF) 内。 |
custom_hooks
以下示例说明了如何为 EMAHook` 提供从 MMPretrain 到 Hydra 配置的自定义 hook
MMPretrain 配置
custom_hooks = [ dict(type='EMAHook', interval=100, priority='HIGH')]
等效的 TAO Hydra 配置
custom_hooks: - type: "EMAHook" momentum: 0.00008 priority: "ABOVE_NORMAL"
有关 custom_hooks 的更多详细信息，请参阅 MMPretrain 文档，了解自定义 hooks。
以下是使用 FAN 主干网络的图像分类 PyT 的模型规范文件示例
model:
backbone:
type: "fan_tiny_8_p4_hybrid"
custom_args:
drop_path_rate: 0.1
freeze: False
pretrained: <Path to pretrained weights>
head:
type: "TAOLinearClsHead"
num_classes: 1000
custom_args:
head_init_scale: 1
loss:
type: LabelSmoothLoss
label_smooth_val: 0.1
mode: 'original'
train_cfg:
augments:
- type: Mixup
alpha: 0.8
- type: CutMix
alpha: 1.0
model 参数主要配置主干网络和 head。
| 参数 | 数据类型 | 默认值 | 描述 | 支持的值 |
|
|
字典 |
无 |
|
|
|
|
字典 |
无 |
包含以下可配置参数 * * 有关基础模型，请参阅 基础模型部分 |
FAN 变体
GCViT 变体
FasterViT 变体 False/True |
head |
字典 | 无 | 分类 head 的配置参数 | – |
train_cfg |
字典 | 无 | 包含高级增强参数。 | – |
基础模型
model:
backbone:
type: "ViT-B-32"
custom_args:
drop_path_rate: 0.1
freeze: False
pretrained: laion400m_e31
head:
type: LinearClsHead
num_classes: 1000
in_channels: 512
loss:
type: CrossEntropyLoss
loss_weight: 1.0
use_soft: False
topk: [1, 5]
受支持的架构和预训练数据集的子集。 请注意，in_channels 应在 head 下更新
CLIP 图像主干网络
| 架构 | 预训练数据集 | in_channels |
|
|
* laion400m_e31 |
512 |
ViT-B-16 |
laion400m_e31 | 512 |
ViT-L-14 |
laion400m_e31 | 768 |
ViT-H-14 |
laion2b_s32b_b79k | 1024 |
ViT-g-14 |
laion2b_s12b_b42k | 1024 |
EVA - CLIP 图像主干网络
| 架构 | 预训练数据集 | in_channels |
EVA02-L-14 |
merged2b_s4b_b131k | 768 |
EVA02-L-14-336 |
laion400m_e31 | 768 |
EVA02-E-14 |
laion400m_e31 | 1024 |
EVA02-E-14-plus |
laion2b_s32b_b79k | 1024 |
head
| 参数 | 数据类型 | 默认值 | 描述 | 支持的值 |
类型 |
字符串 | 无 | 用于生成混合比率的 Beta 分布参数 | TAOLinearClsHead, FANLinearClsHead, LogisticRegressionHead |
num_classes |
字典 | 无 | 训练类的数量 | >=0 |
loss |
字典 | {“type”:”CrossEntropyLoss”} | 有关不同类型的损失及其参数，请参阅 losses | |
topk |
列表 | [1,] | 类的数量 | >=0 |
|
|
字典 |
无 |
要传递给 |
– |
|
|
字典 |
无 |
用于 Logistic Regression Head 的参数 |
– |
lr_head
Logistic Regression head 由以下参数定义
|
|
字典 |
无 |
包含以下可调参数 |
– |
Logistic Regression head 涉及冻结主干网络权重。 请注意，freeze 参数应在主干网络配置下更新为 True，如下所示
model:
backbone:
type: "ViT-B-32"
custom_args:
drop_path_rate: 0.1
freeze: True
pretrained: laion400m_e31
head:
type: "LogisticRegressionHead"
num_classes: 1000
lr_head:
C: 0.316
max_iter: 5000
train_cfg
Mixup
| 参数 | 数据类型 | 默认值 | 描述 | 支持的值 |
alpha |
字符串 | 无 | 用于生成混合比率的 Beta 分布参数 | 0-1 |
CutMix
| 参数 | 数据类型 | 默认值 | 描述 | 支持的值 |
alpha |
字符串 | 无 | 用于生成混合比率的 Beta 分布参数 | 0-1 |
cutmix_minmax |
list[float] | 无 | 补丁的最小/最大面积比率。 | |
correct_lam |
布尔值 | True | 当 cutmix bbox 被图像边界裁剪时，是否应用 lambda 校正。 | False/True |
loss
下面显示了一些重要的分类损失。 请注意，通过遵循 TAO 的 Hydra 配置，可以使用 MMPretrain 中所有支持的损失。 有关 MMPretrain 损失的列表，请参阅 losses_mmpretrain 文档。
LabelSmoothLoss
| 参数 | 数据类型 | 默认值 | 描述 | 支持的值 |
label_smooth_val |
字符串 | 无 | 标签平滑的程度 | 0-1 |
use_sigmoid |
布尔值 | 无 | 指定预测是否应使用 sigmoid 或 softmax | False/ True |
num_classes |
整数 | 无 | 类的数量 | >=0 |
mode |
字符串 | 无 | 用于生成混合比率的 Beta 分布参数 | 0-1 |
reduction |
字符串 | 无 | 用于减少损失的方法 | 均值, 总和 |
loss_weight |
浮点数 | 1.0 | 损失的权重 | >=0 |
CrossEntropyLoss
| 参数 | 数据类型 | 默认值 | 描述 | 支持的值 |
use_sigmoid |
布尔值 | False | 指定预测是否应使用 sigmoid 或 softmax | 0-1 |
use_soft |
布尔值 | False | 指定是否使用 CrossEntropyLoss 的 soft 版本 | 0-1 |
loss_weight |
浮点数 | 1.0 | 损失的权重 | 0-1 |
reduction |
字符串 | 无 | 用于减少损失的方法 | 均值, 总和 |
class_weight |
list[float] | 无 | 每个类的权重 | 0-1 |
使用 tao model classification_pyt train 命令训练分类 pytorch 模型
tao model classification_pyt 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。
evaluate 配置定义评估过程的超参数。 以下是配置示例
evaluate:
checkpoint: /path/to/model.pth
topk: 1
在使用实验配置文件并按照训练模型的步骤训练模型后，下一步是在测试集上评估此模型，以衡量模型的准确性。 TAO 包括 tao model classification_pyt evaluate 命令来执行此操作。
分类应用程序计算评估损失和 Top-k 准确率。
训练后，模型存储在您选择的输出目录 results_dir 中。
evaluate:
checkpoint: /path/to/model.pth
tao model classification_pyt evaluate [-h] -e <experiment_spec>
evaluate.checkpoint=<model to be evaluated>
results_dir=<path to results dir>
[evaluate.<evaluate_option>=<evaluate_option_value>]
[evaluate.gpu_ids=<gpu indices>]
[evaluate.num_gpus=<number of gpus>]
必需参数
-e, --experiment_spec: 用于设置评估实验的实验规范文件evaluate.checkpoint: 要评估的.pth模型。results_dir: 结果将存储在其中的路径
可选参数
evaluate.<evaluate_option>: 评估选项。
对于分类，tao model classification_pyt inference 保存一个 .csv 文件，其中包含多个图像的图像路径和相应的标签。 也可以启用 TensorRT Python 推理。
inference:
checkpoint: /path/to/model.pth
tao model classification_pyt inference [-h] -e <experiment_spec_file>
inference.checkpoint=<model to be inferenced>
results_dir=<path to results dir>
[inference.<inference_option>=<inference_option_value>]
[inference.gpu_ids=<gpu indices>]
[inference.num_gpus=<number of gpus>]
必需参数
-e, --experiment_spec: 用于设置推理实验的实验规范文件inference.checkpoint: 要推理的.pth模型。results_dir: 结果将存储在其中的路径
可选参数
inference.<inference_option>: 推理选项。
导出模型将训练过程与推理分离，并允许在 TAO 环境之外转换为 TensorRT 引擎。 TensorRT 引擎特定于每个硬件配置，应为每个唯一的推理环境生成。 导出的模型可以在训练和部署硬件之间通用使用。 导出的模型格式称为 .onnx。
export 参数定义导出过程的超参数。
export:
checkpoint: /path/to/model.pth
onnx_file: /path/to/model.onnx
opset_version: 12
verify: False
input_channel: 3
input_width: 224
input_height: 224
以下是 tao classification_pyt export 命令的示例
tao model classification_pyt export [-h] -e <experiment spec file>
export.checkpoint=<model to export>
export.onnx_file=<onnx path>
[export.<export_option>=<export_option_value>]
必需参数
-e, --experiment_spec: 实验规范文件的路径export.checkpoint: 要导出的.pth模型。export.onnx_file: 保存.etlt或.onnx模型的路径。
可选参数
export.<export_option>: 导出选项。
有关 TensorRT 引擎生成、验证和 INT8 校准，请参阅 TAO Deploy 文档。
有关使用 DeepStream 部署分类模型的更多信息，请参阅 集成分类 (TF1/TF2/PyTorch) 模型页面。