度量学习识别 - NVIDIA 文档

度量学习识别 (MLRecogNet) 是一种分类器，可将输入图像编码为嵌入向量，并根据参考空间中的嵌入向量预测其标签。MLRecogNet 由两部分组成：

Trunk：一个骨干网络，将输入图像编码为特征向量。
Embedder：一个全连接层，将特征向量映射到嵌入空间。

嵌入空间是一个高维空间，其中同一类别的嵌入向量之间的距离很小，不同类别的嵌入向量之间的距离很大。Embedder 经过训练，可以最大限度地缩小同一类别的嵌入向量之间的距离，并最大限度地扩大不同类别的嵌入向量之间的距离。查询图像的嵌入向量与参考图像的嵌入向量进行比较，以预测查询图像的标签。

当前支持的 trunk 是 ResNet，它是视觉分类最常用的基线。当前支持的 embedder 是单层 MLP。

在训练、评估和推理期间，MLRecogNet 需要参考集和查询集以进行验证或测试。参考集由标记图像的集合组成，而查询集指的是一组未标记的图像——目标是通过比较未标记图像与训练后的 MLRecogNet 生成的参考集的嵌入向量的相似性来预测未标记图像的标签。

准备数据集

MLRecogNet 需要来自检测集或分类集的裁剪图像作为输入。默认情况下，这些图像会被调整大小为 224x224 以作为模型输入。在训练期间，会向每个图像应用增强。

数据应按以下结构组织：

复制
已复制！

            
            /Dataset
    /reference
      /class1
            0001.jpg
            0002.jpg
            ...
            0100.jpg
      /class2
            0001.jpg
            0002.jpg
            ...
            0100.jpg
        ...
    /train
      /class1
            0101.jpg
            0102.jpg
            ...
            0200.jpg
      /class2
            0101.jpg
            0102.jpg
            ...
            0200.jpg
    /val
      /class1
            0201.jpg
            0202.jpg
            ...
            0220.jpg
      /class2
            0201.jpg
            0202.jpg
            ...
            0220.jpg
    /test
      /class1
            0301.jpg
            0302.jpg
            ...
            0400.jpg
      /class2
            0301.jpg
            0302.jpg
            ...
            0400.jpg

数据集的根目录包含用于参考、训练、验证和测试的子目录。子目录需要采用 ImageNet 结构，如上所示。每个子目录都包含同一类别的图像。如果 test 集中的类别不在 reference 集中，则无法正确识别查询的图像。

创建实验规范文件

MLRecogNet 的规范文件包括 model、train 和 dataset 参数。以下是用于在目标数据集上训练 MLRecogNet 模型的示例规范 $TRAIN_SPEC。

复制
已复制！

            
            results_dir: "???"
model:
  backbone: resnet_101
  pretrained_model_path: /path/to/resnet101_pretrained_mlrecog.pth.tar
  input_width: 224
  input_height: 224
  feat_dim: 2048
train:
  optim:
    name: Adam
    steps: [40, 70]
    gamma: 0.1
    embedder:
      bias_lr_factor: 1
      weight_decay: 0.0001
      weight_decay_bias: 0.0005
      base_lr: 0.000001
      momentum: 0.9
    trunk:
      bias_lr_factor: 1
      weight_decay: 0.0001
      weight_decay_bias: 0.0005
      base_lr: 0.00001
      momentum: 0.9
    warmup_factor: 0.01
    warmup_iters: 10
    warmup_method: linear
    triplet_loss_margin: 0.3
    miner_function_margin: 0.1
  num_epochs: 10
  resume_training_checkpoint_path: null
  checkpoint_interval: 5
  validation_interval: 5
  smooth_loss: False
  batch_size: 16
  val_batch_size: 16
  seed: 1234
dataset:
  train_dataset: /path/to/dataset/train
  val_dataset:
    reference: /path/to/dataset/reference
    query: /path/to/dataset/val
  workers: 12
  pixel_mean: [0.485, 0.456, 0.406]
  pixel_std: [0.226, 0.226, 0.226]
  prob: 0.5
  re_prob: 0.5
  num_instance: 4
  color_augmentation:
    enabled: True
    brightness: 0.5
    contrast: 0.3
    saturation: 0.1
    hue: 0.1
  gaussian_blur:
    enabled: True
    kernel: [15, 15]
    sigma: [0.3, 0.7]
  random_rotation: True
  class_map: /path/to/class_map.yaml

参数	数据类型	默认值	描述	支持的值
`model`	dict config	–	模型架构的配置
`dataset`	dict config	–	数据集的配置
`train`	dict config	–	训练任务的配置
`evaluate`	dict config	–	评估任务的配置
`inference`	dict config	–	推理任务的配置
`encryption_key`	string	None	用于加密和解密模型文件的加密密钥
`results_dir`	string	/results	保存实验结果的目录
`export`	dict config	–	ONNX 导出任务的配置
`gen_trt_engine`	dict config	–	TensorRT 生成任务的配置。仅在 tao deploy 中使用

model

model 参数提供更改 MetricLearningRecognition 架构的选项。

复制
已复制！

            
            model:
  backbone: resnet_50
  pretrained_model_path: "/path/to/pretrained_model.pth"
  pretrained_embedder_path: null
  pretrained_trunk_path: null
  input_channels: 3
  input_width: 224
  input_height: 224
  feat_dim: 256

参数	数据类型	默认值	描述	支持的值
`backbone`	string	resnet_50	骨干 (trunk) 模型类型。	resnet_50, resnet_101, fan_small, fan_base, fan_large, fan_tiny, nvdinov2_vit_large_legacy
`pretrained_model_path`	string		预训练模型的路径。权重仅加载到完整模型
`pretrained_trunk_path`	string		预训练 trunk 的路径。权重仅加载到 trunk 部分。
`pretrained_embedder_path`	string		预训练 embedder 的路径。权重仅加载到 embedder 部分。
`input_channels`	unsigned int	3	输入通道数	>0
`input_width`	int	224	图像的输入宽度	int
`input_height`	int	224	图像的输入高度	int
`feat_dim`	unsigned int	256	特征嵌入的输出大小	>0

train

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

复制
已复制！

            
            train:
  optim:
    name: Adam
    steps: [40, 70]
    gamma: 0.1
    warmup_factor: 0.01
    warmup_iters: 10
    warmup_method: 'linear'
    triplet_loss_margin: 0.3
    miner_function_margin: 0.1
    embedder:
      bias_lr_factor: 1
      base_lr: 0.000001
      momentum: 0.9
      weight_decay: 0.0001
      weight_decay_bias: 0.0005
    trunk:
      bias_lr_factor: 1
      base_lr: 0.00001
      momentum: 0.9
      weight_decay: 0.0001
      weight_decay_bias: 0.0005
  num_epochs: 10
  checkpoint_interval: 5
  validation_interval: 5
  clip_grad_norm: 0.0
  resume_training_checkpoint_path: null
  report_accuracy_per_class: True
  smooth_loss: True
  batch_size: 64
  val_batch_size: 64
  train_trunk: false
  train_embedder: true
  results_dir: null
  seed: 1234

参数	数据类型	默认值	描述	支持的值
`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	保存训练结果的目录
`optim`	dict config	–	torch 优化器（Optim Config）的配置，包括学习率、学习率调度器、权重衰减等。
`clip_grad_norm`	float	0.0	通过 L2 范数裁剪梯度的量。值为 0.0 表示不裁剪。	>=0
`report_accuracy_per_class`	bool	True	如果为 True，则将报告每个类别的 top1 精度。	True/False
`smooth_loss`	bool	True	如果为 True，将使用 triplet 损失的 log-exp 版本。	True/False
`batch_size`	unsigned int	64	训练的批次大小	>0
`val_batch_size`	unsigned int	64	验证的批次大小	>0
`train_trunk`	bool	True	如果为 False，则在训练期间将冻结模型的 trunk 部分	True/False
`train_embedder`	bool	True	如果为 False，则在训练期间将冻结模型的 embedder 部分	True/False

optim

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

复制
已复制！

            
            optim:
  name: Adam
  steps: [40, 70]
  gamma: 0.1
  warmup_factor: 0.01
  warmup_iters: 10
  warmup_method: 'linear'
  triplet_loss_margin: 0.3
  miner_epsilon: 0.1
  embedder:
    bias_lr_factor: 1
    base_lr: 0.00035
    momentum: 0.9
    weight_decay: 0.0005
    weight_decay_bias: 0.0005
  trunk:
    bias_lr_factor: 1
    base_lr: 0.00035
    momentum: 0.9
    weight_decay: 0.0005
    weight_decay_bias: 0.0005

参数	数据类型	默认值	描述	支持的值
`name`	string	Adam	优化器的名称。支持 `torch.optim` 中的算法。	Adam/SGD/Adamax/…
`steps`	int list	[40, 70]	用于 `MultiStep` 调度器降低学习率的步数
`gamma`	float	0.1	`WarmupMultiStepLR` 调度器的衰减率	>0.0
`warmup_factor`	float	0.01	`WarmupMultiStepLR` 调度器的预热因子	>0.0
`warmup_iters`	unsigned int	10	`WarmupMultiStepLR` 调度器的预热迭代次数	>0
`warmup_method`	string	linear	优化器的预热方法	constant/linear
`triplet_loss_margin`	float	0.3	锚点-正样本距离与锚点-负样本距离之间所需的差异	>0.0
`miner_function_margin`	float	0.1	如果负样本对的相似度大于最难正样本对的相似度减去此边距，则选择负样本对；如果正样本对的相似度小于最难负样本对的相似度加上此边距，则选择正样本对	>0.0
`embedder`	dict config	–	MLRecogNet embedder 的学习率配置 (LR Config)
`trunk`	dict config	–	MLRecogNet trunk 的学习率配置 (LR Config)

LR Config

参数	数据类型	默认值	描述	支持的值
`base_lr`	float	0.00035	训练的初始学习率	>0.0
`bias_lr_factor`	float	1	WarmupMultiStepLR 的偏差学习率因子	>=1
`momentum`	float	0.9	WarmupMultiStepLR 优化器的动量	>0.0
`weight_decay`	float	0.0005	优化器的权重衰减系数	>0.0
`weight_decay_bias`	float	0.0005	优化器的权重衰减偏差	>0.0

dataset

dataset 参数定义数据集源、训练批次大小和增强。

复制
已复制！

            
            dataset:
  train_dataset: /path/to/dataset/train
  val_dataset:
    reference: /path/to/dataset/reference
    query: /path/to/dataset/val
  workers: 8
  pixel_mean: [0.485, 0.456, 0.406]
  pixel_std: [0.226, 0.226, 0.226]
  padding: 10
  prob: 0.5
  re_prob: 0.5
  sampler: softmax_triplet
  num_instance: 4
  gaussian_blur:
    enabled: True
    kernel: [15, 15]
    sigma: [0.3, 0.7]
  color_augmentation:
    enabled: True
    brightness: 0.5
    contrast: 0.3
    saturation: 0.1
    hue: 0.1

参数	数据类型	默认值	描述	支持的值
`train_dataset`	string		训练数据集的路径。此字段仅在 train 任务中是必需的。
`val_dataset`	dict		参考集和查询集地址的映射。对于训练和评估，这两个字段都是必需的。对于推理，仅需要参考集地址。	{“reference”: /path/to/reference/set, “query”: “”}
`workers`	unsigned int	8	并行处理数据的工作线程数	>0
`class_map`	string		在 `tao model` 中，class_map 是一个 YAML 文件，用于将数据集类名称映射到所需的类名称。如果未指定，默认情况下，报告的类名称是数据集文件夹中的文件夹名称。在 `tao deploy` 中，class_map 是一个 TXT 文件，逐行列出类名称。行索引将是类索引。默认情况下，类名称是文件夹名称，并且类的顺序是字母数字顺序。
`pixel_mean`	float list	[0.485, 0.456, 0.406]	用于图像归一化的像素均值	float list
`pixel_std`	float list	[0.226, 0.226, 0.226]	用于图像归一化的像素标准差	float list
`num_instance`	unsigned int	4	一个批次中同一人的图像实例数	>0
`prob`	float	0.5	用于图像增强的随机水平翻转概率	>0
`re_prob`	float	0.5	用于图像增强的随机擦除概率	>0
`random_rotation`	bool	True	如果为 True，则对输入数据应用 0 ~ 180 度的随机旋转	True/False
`gaussian_blur`	dict config	–	输入样本上的高斯模糊增强的配置
`color_augmentation`	dict config	–	输入样本上的颜色增强的配置

高斯模糊配置

参数	数据类型	默认值	描述	支持的值
`enabled`	bool	True	如果为 True，则对输入样本应用高斯模糊增强	True/False
`kernel`	unsigned int list	[15, 15]	高斯模糊的内核大小
`sigma`	float list	[0.3, 0.7]	高斯模糊的 sigma 值范围

颜色增强配置

参数	数据类型	默认值	描述	支持的值
`enabled`	bool	True	如果为 True，则对输入样本应用颜色增强	True/False
`brightness`	float	0.5	抖动亮度的值	>=0
`contrast`	float	0.3	抖动对比度的值	>=0
`saturation`	float	0.1	抖动饱和度的值	>=0
`hue`	float	0.1	抖动色调的值	>=0, <=0.5

训练模型

使用以下命令运行 MLRecogNet 训练

复制
已复制！

            
            tao model ml_recog train -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.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'

最新的检查点也保存为 ml_model_latest.pth。如果 ml_model_latest.pth 存在于 train.results_dir 中，则训练会自动从该文件恢复。如果提供了 train.resume_training_checkpoint_path，则会取代此行为。

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

以下是输出 $RESULTS_DIR/train/status.json 的示例

复制
已复制！

            
            {"date": "6/20/2023", "time": "23:11:2", "status": "STARTED", "verbosity": "INFO", "message": "Starting Training Loop."}
...

{"date": "6/20/2023", "time": "23:11:22", "status": "SUCCESS", "verbosity": "INFO", "message": "Train finished successfully."}

评估模型

以下是用于在测试数据集上评估 MLRecogNet 模型的示例规范 $EVAL_SPEC。

复制
已复制！

            
            results_dir: /path/to/root/results/dir
model:
  backbone: resnet_50
  input_width: 224
  input_height: 224
  feat_dim: 256
dataset:
  workers: 8
  val_dataset:
    reference: /path/to/dataset/reference
    query: /path/to/dataset/val
evaluate:
  checkpoint: /path/to/checkpoint
  batch_size: 128
  results_dir: /path/to/results

参数	数据类型	默认值	描述	支持的值
`checkpoint`	string	None	要评估的 .pth Torch 模型的路径
`results_dir`	string	/results/evaluate	保存评估结果的目录
`num_gpus`	unsigned int	1	用于分布式评估的 GPU 数量	>0
`gpu_ids`	List[int]	[0]	用于分布式评估的 GPU 索引
`trt_engine`	string	None	要评估的 TensorRT (TRT) 引擎的路径。目前，TAO Deploy 仅支持 `trt_engine`
`topk`	int	1	如果大于 1，则精度将为 top-k 精度。目前，TAO Deploy 仅支持 `evaluate.topk`	>0
`batch_size`	int	64	评估任务的批次大小	>0
`report_accuracy_per_class`	bool	True	如果为 True，则将报告每个类别的 top-1 精度	True/False

以下是 MLRecogNet 的评估指标：

Adjusted Mutual Information (AMI)：统计学和信息论中使用的一种度量，用于量化两个分配之间的一致性，例如聚类分配，该度量针对偶然性进行了调整，因此与原始互信息相比，可以更准确地描述两者之间的相似性。
Normalized Mutual Information (NMI)：互信息 (MI) 分数的归一化，用于将结果缩放到 0（无互信息）和 1（完全相关）之间。
Mean Average Precision：模型在不同召回率水平下实现的平均精度，全面评估其在信息检索方面的性能。
Mean Average Precision at r：模型针对前 R 个排名结果的平均精度，当考虑有限数量的结果时，可以深入了解模型的检索或对象检测性能的有效性。
Mean Reciprocal Rank：一组查询的第一个相关结果的倒数排名的平均值，强调尽早检索相关信息的重要性。
Precision at 1：最近邻检索的准确率。
R Precision：信息检索系统的评估指标，用于衡量前 R 个排名结果中相关文档的比例，其中“R”对应于给定查询的相关文档总数。

当 evaluate.report_accuracy_per_class 设置为 True 时，将添加每个类别的准确率。

使用以下命令运行 MLRecogNet 评估

复制
已复制！

            
            tao model ml_recog evaluate -e <experiment_spec_file>
                   evaluate.checkpoint=<model to be evaluated>
                   dataset.val_dataset.reference=<path to test reference set>
                   dataset.val_dataset.query=<path to test query set>
                   [evaluate.<evaluate_option>=<evaluate_option_value>]
                   [evaluate.gpu_ids=<gpu indices>]
                   [evaluate.num_gpus=<number of gpus>]

必需参数

-e, --experiment_spec_file：用于设置评估实验的实验规范文件
evaluate.checkpoint：要评估的 .pth 模型的路径
dataset.val_dataset.reference：测试参考集的路径
dataset.val_dataset.query：测试查询集的路径

可选参数

evaluate.<evaluate_option>：评估选项。

以下是输出 $RESULTS_DIR/evaluate/status.json 的示例

复制
已复制！

            
            {"date": "6/2/2023", "time": "6:12:16", "status": "STARTED", "verbosity": "INFO", "message": "Starting Metric Learning Recognition evaluate."}
{"date": "6/2/2023", "time": "6:12:17", "status": "STARTED", "verbosity": "INFO", "message": "Loading checkpoint:$RESULTS_DIR/train/ml_model_epoch=000.pth"}
{"date": "6/2/2023", "time": "6:12:17", "status": "RUNNING", "verbosity": "INFO", "message": "Constructing model graph..."}
{"date": "6/2/2023", "time": "6:12:17", "status": "SKIPPED", "verbosity": "INFO", "message": "Skipped loading pretrained model as checkpoint is to load."}
{"date": "6/2/2023", "time": "6:12:23", "status": "SUCCESS", "verbosity": "INFO", "message": "Evaluate finished successfully.", "kpi": {"AMI": 0.8074901483322209, "NMI": 0.8118350536509751, "Mean Average Precision": 0.6876838920302153, "Mean Reciprocal Rank": 0.992727267742157, "r-Precision": 0.666027864375903, "Precision at Rank 1": 0.989090909090909}}

以下是打印输出的示例

复制
已复制！

            
            Starting Metric Learning Recognition evaluate.
Experiment configuration:
...
results_dir: $RESULTS_DIR

Loading checkpoint: $RESULTS_DIR/train/ml_model_epoch=000.pth
Constructing model graph...
Skipped loading pretrained model as checkpoint is to load.
Evaluating epoch eval mode
...
Computing accuracy for the query split w.r.t ['gallery']
running k-nn with k=106
embedding dimensionality is 256
/usr/local/lib/python3.8/dist-packages/torch/storage.py:315: UserWarning: TypedStorage is deprecated. It will be removed in the future and UntypedStorage will be the only storage class. This should only matter to you if you are using storages directly.
  warnings.warn(message, UserWarning)
running k-means clustering with k=5
embedding dimensionality is 256
******************* Evaluation results **********************
AMI: 0.8075
NMI: 0.8118
Mean Average Precision: 0.7560
Mean Reciprocal Rank: 0.9922
r-Precision: 0.7421
Precision at Rank 1: 0.9882
*************************************************************

在模型上运行推理

以下是用于在推理集上运行 MLRecogNet 模型推理的示例规范 $INFERENCE_SPEC

复制
已复制！

            
            results_dir: /path/to/root/results/dir
model:
  backbone: resnet_50
  input_width: 224
  input_height: 224
  feat_dim: 256
dataset:
  workers: 8
  val_dataset:
    reference: /path/to/dataset/reference
    query: ""
inference:
  input_path: /path/to/dataset/test
  inference_input_type: classification_folder
  checkpoint: /path/to/model/checkpoint
  results_dir: /path/to/results/dir
  batch_size: 128

参数	数据类型	默认值	描述	支持的值
`checkpoint`	string	None	要运行推理的 .pth torch 模型的路径
`results_dir`	string	/results/inference	保存推理结果的目录
`num_gpus`	unsigned int	1	用于分布式推理的 GPU 数量	>0
`gpu_ids`	List[int]	[0]	用于分布式推理的 GPU 索引
`trt_engine`	string	None	要运行推理的 TensorRT (TRT) 引擎的路径。目前，TAO Deploy 仅支持 `trt_engine`。
`input_path`	string		要对其运行推理的数据的路径	>0
`inference_input_type`	string	“image_folder”	支持三个选项 `image_folder`：当 `input_path` 是图像文件夹时使用。 `classification_folder`：当 `input_path` 是 `ImageNet` 结构化文件夹时使用。 `image`：当 `input_path` 是图像文件时使用	“image_folder”/”classification_folder”/”image”
`batch_size`	int	64	推理任务的批次大小	>0
`topk`	int	1	要返回的 top 结果数	>0

使用以下命令在 MLRecogNet 上使用 .pth 模型运行推理

复制
已复制！

            
            tao model ml_recog inference -e <experiment_spec>
                                inference.checkpoint=<inference model>
                                dataset.val_dataset.reference=<path to gallery data>
                                inference.input_path=<path to query data>
                                [inference.<inference_option>=<inference_option_value>]
                                [inference.gpu_ids=<gpu indices>]
                                [inference.num_gpus=<number of gpus>]

输出是一个 CSV 文件，其中包含所有查询数据的特征嵌入及其预测标签。

必需参数

-e, --experiment_spec：用于设置推理的实验规范文件
inference.checkpoint：用于执行推理的 .pth 模型
dataset.val_dataset.reference：参考集的路径
inference.input_path：要对其运行推理的数据的路径

可选参数

inference.<inference_option>：推理选项。

预期输出如下：

复制
已复制！

            
            /path/to/images/c000001_10.png,"['c000001', 'c000005', 'c000001', 'c000005']","[5.0030694183078595e-06, 5.5495906963187736e-06, 5.976316515443614e-06, 6.004379429214168e-06]"
/path/to/images/c000001_11.png,"['c000001', 'c000005', 'c000001', 'c000001']","[3.968068540416425e-06, 5.043690180173144e-06, 5.885293830942828e-06, 6.030047643434955e-06]"
/path/to/images/c000001_120.png,"['c000001', 'c000001', 'c000005', 'c000003']","[1.9612791675172048e-06, 4.112744136364199e-06, 4.603011802828405e-06, 5.8091877690458205e-06]"

其中第一列包含推理图像路径，第二列包含 top-k 预测标签，第三列包含 top-k 结果的嵌入向量距离。

以下是输出 $RESULTS_DIR/inference/status.json 的示例

复制
已复制！

            
            {"date": "6/2/2023", "time": "6:13:47", "status": "STARTED", "verbosity": "INFO", "message": "Starting Metric Learning Recognition inference."}
{"date": "6/2/2023", "time": "6:13:47", "status": "STARTED", "verbosity": "INFO", "message": "Loading checkpoint:$RESULTS_DIR/train/ml_model_epoch=001.pth"}
{"date": "6/2/2023", "time": "6:13:47", "status": "RUNNING", "verbosity": "INFO", "message": "Constructing model graph..."}
{"date": "6/2/2023", "time": "6:13:48", "status": "SKIPPED", "verbosity": "INFO", "message": "Skipped loading pretrained model as checkpoint is to load."}
{"date": "6/2/2023", "time": "6:14:6", "status": "SUCCESS", "verbosity": "INFO", "message": "result saved at$RESULTS_DIR/inference/result.csv"}
{"date": "6/2/2023", "time": "6:14:6", "status": "SUCCESS", "verbosity": "INFO", "message": "Inference finished successfully."}

以下是打印输出的示例

复制
已复制！

            
            Starting Metric Learning Recognition inference.
Experiment configuration:
...

Loading checkpoint: $RESULTS_DIR/train/ml_model_epoch=001.pth
Constructing model graph...
Skipped loading pretrained model as checkpoint is to load.
/usr/local/lib/python3.8/dist-packages/torch/storage.py:315: UserWarning: TypedStorage is deprecated. It will be removed in the future and UntypedStorage will be the only storage class. This should only matter to you if you are using storages directly.
  warnings.warn(message, UserWarning)
...
result saved at $RESULTS_DIR/inference/result.csv
Inference finished successfully.

导出模型

以下是用于导出 MLRecogNet 模型的示例规范 $EXPORT_SPEC

复制
已复制！

            
            results_dir: /path/to/root/results/dir
model:
  backbone: resnet_50
  input_width: 224
  input_height: 224
  feat_dim: 256
export:
  checkpoint: /path/to/checkpoint
  onnx_file: /path/to/results/model.onnx
  results_dir:  /path/to/results
  batch_size: -1
  on_cpu: false
  verbose: true

参数	数据类型	默认值	描述	支持的值
`checkpoint`	string	None	要评估的 .pth Torch 模型的路径
`onnx_file`	string	None	导出的 ONNX 文件的路径。如果未指定此值，则默认为 `export.results_dir` 中的 `model.onnx`
`batch_size`	int	-1	导出的 ONNX 模型的批次大小。如果 `batch_size` 为 -1，则导出的 ONNX 模型具有动态批次大小。	>0; -1
`gpu_id`	unsigned int	0	用于 Torch 到 ONNX 导出的 GPU ID。目前，导出任务仅支持在单个 GPU 上运行	>=0
`on_cpu`	bool	False	如果为 True，则将在 CPU 上执行 Torch 到 ONNX 导出	True/False
`opset_version`	unsigned int	14	要定位的默认 (ai.onnx) opset 的版本	>= 7 且 <= 16。
`verbose`	bool	True	如果为 True，则将要导出的模型的描述打印到 `stdout`。	True/False
`results_dir`	string	None	导出任务的结果目录的路径

使用以下命令将 MLRecogNet 导出为 .onnx 格式以进行部署

复制
已复制！

            
            tao model ml_recog export -e <experiment_spec>
                          export.checkpoint=<.pth checkpoint to be exported>
                          [export.onnx_file=<path to exported ONNX file>]
                          [export.<export_option>=<export_option_value>]

必需参数

-e, --experiment_spec：用于设置导出的实验规范文件
export.checkpoint：要导出的 .pth 模型

可选参数

export.onnx_file：用于保存导出模型的路径。默认路径与 export.results_dir（如果有）或 results_dir 位于同一目录中。
export.<export_option>：导出选项。

以下是输出 $RESULTS_DIR/export/status.json 的示例

复制
已复制！

            
            {"date": "6/2/2023", "time": "6:17:45", "status": "STARTED", "verbosity": "INFO", "message": "Starting Metric Learning Recognition export."}
{"date": "6/2/2023", "time": "6:17:45", "status": "STARTED", "verbosity": "INFO", "message": "Loading checkpoint:$RESULTS_DIR/train/ml_model_epoch=001.pth"}
{"date": "6/2/2023", "time": "6:17:45", "status": "RUNNING", "verbosity": "INFO", "message": "Constructing model graph..."}
{"date": "6/2/2023", "time": "6:17:46", "status": "SKIPPED", "verbosity": "INFO", "message": "Skipped loading pretrained model as checkpoint is to load."}
{"date": "6/2/2023", "time": "6:17:46", "status": "STARTED", "verbosity": "INFO", "message": "Exporting model to ONNX"}
{"date": "6/2/2023", "time": "6:17:48", "status": "STARTED", "verbosity": "INFO", "message": "Simplifying ONNX model"}
{"date": "6/2/2023", "time": "6:17:50", "status": "SUCCESS", "verbosity": "INFO", "message": "ONNX model saved at$RESULTS_DIR/export/ml_model_epoch=001.onnx"}
{"date": "6/2/2023", "time": "6:17:50", "status": "SUCCESS", "verbosity": "INFO", "message": "Export finished successfully."}

以下是打印输出的示例

复制
已复制！

            
            Starting Metric Learning Recognition export.
Experiment configuration:
...

Loading checkpoint: $RESULTS_DIR/train/ml_model_epoch=001.pth
Constructing model graph...
Skipped loading pretrained model as checkpoint is to load.
Exporting model to ONNX
Exported graph: graph(%input : Float(*, 3, 224, 224, strides=[150528, 50176, 224, 1], requires_grad=0, device=cuda:0),
  ...

========== Diagnostic Run torch.onnx.export version 1.14.0a0+44dac51 ===========
verbose: False, log level: Level.ERROR
======================= 0 NONE 0 NOTE 0 WARNING 0 ERROR ========================

Simplifying ONNX model
Checking 0/3...
Checking 1/3...
Checking 2/3...
ONNX model saved at $RESULTS_DIR/export/ml_model_epoch=001.onnx
Export finished successfully.

部署模型

您可以使用 TAO Deploy 在边缘设备（例如 Jetson Xavier、Jetson Nano 或 Tesla）或云中 NVIDIA GPU 上部署训练好的深度学习和计算机视觉模型。TAO Deploy 是 TAO 中的一个应用程序，可将 ONNX 模型转换为 TensorRT 引擎，并通过 TensorRT 引擎运行推理。

在 TAO Deploy 上运行 MLRecogNet 推理

从导出生成的 MLRecogNet ONNX 文件作为 TAO Deploy 的输入，以生成优化的 TensorRT 引擎。有关使用 TAO Deploy 在 MLRecogNet TensorRT 引擎上运行推理的更多信息，请参阅 TAO Deploy 文档。