目录
1. 简介
2. 代码详解
2.1 导入所需的库
2.2 创建 Inspector
2.3 下载模型
2.4 检查模型
3. 其他有用函数
3.1 查看 torchvision 中模型
3.2 保存模型
3.2.1 保存模型参数
3.2.2 保存完整模型
3.2.3 加载模型
4. 总结
1. 简介
在《Vitis AI 构建开发环境,使用inspector检查模型_vitis-ai 开发-CSDN博客》一文中,笔者简单分享了 inspector 的简单用法,本文将深入探究其用法。
inspector 的函数,用来诊断不同器件架构下的神经网络 (NN) 模型。检查器可以基于硬件约束来预测目标器件分配。生成的检查报告可用于指导用户对 NN 模型进行修改或最优化,从而显著降低部署难度并缩短部署时间。
建议在量化浮点模型前对其进行检查。
NNDCT:Neural Network Development and Compiler Tools。
2. 代码详解
2.1 导入所需的库
Vitis AI 2.5, Pytorch 版本信息:
Python 3.7.12
PyTorch 1.10.1
torchvision 0.11.2
import torch
from torchvision.models import resnet18
from pytorch_nndct.apis import Inspector
from IPython.display import Image
功能解释:
- import torch:PyTorch 库,用于深度学习的开源框架,核心功能是张量计算(类似于NumPy)以及深度神经网络的自动求导机制。
- torchvision.models:建立在 PyTorch 之上的一个库,专门用来处理图像数据。提供了三大主要功能。
- 一是加载和预处理图像数据的工具。
- 二是常用的图像数据集,如ImageNet、CIFAR10、MNIST等。
- 三是预训练好的模型,如VGG、ResNet等。
- Inspector 是 Vitis AI 工具的一部分,用于检查和诊断浮点模型。
2.2 创建 Inspector
target = "DPUCZDX8G_ISA1_B4096"
inspector = Inspector(target)
---
#输出
---
[VAIQ_NOTE]: =>Inspector is initialized successfully with target:
name: DPUCZDX8G_ISA1_B4096
type: DPUCZDX8G
isa_version: 1
该方法来自于类:
# <Vitis-AI-2.5/src/Vitis-AI-Quantizer/vai_q_pytorch/pytorch_binding/pytorch_nndct/apis.py>
---
class Inspector(object):
def __init__(self, name_or_fingerprint: str):
两种创建方法:
inspector = Inspector("0x603000b16013831") # by target fingerprint
# or
inspector = Inspector("DPUCZDX8G_ISA1_B4096") # by target name
2.3 下载模型
model = resnet18(pretrained=True)
model.eval()
dummy_input = torch.randn(1, 3, 224, 224)
2.4 检查模型
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
inspector.inspect(model, (dummy_input,), device=device, output_dir="inspect", image_format="png")
- device:指定模型在检查过程中的计算设备。
检查完毕,输出:
[VAIQ_NOTE]: =>Start to inspect model...
[VAIQ_NOTE]: =>Quant Module is in 'cpu'.
[VAIQ_NOTE]: =>Parsing ResNet...
[VAIQ_NOTE]: Start to trace model...
[VAIQ_NOTE]: Finish tracing.
[VAIQ_NOTE]: Processing ops...
██████████████████████████████████████████████████| 71/71 [00:00<00:00, 750.15it/s, OpInfo: name = return_0, type = Return]
[VAIQ_NOTE]: =>Doing weights equalization...
[VAIQ_NOTE]: =>Quantizable module is generated.(inspect/ResNet.py)
[VAIQ_NOTE]: All the operators are assigned to the DPU(see more details in 'inspect/inspect_DPUCZDX8G_ISA1_B4096.txt')
[VAIQ_NOTE]: Dot image is generated.(inspect/inspect_DPUCZDX8G_ISA1_B4096.png)
[VAIQ_NOTE]: =>Finish inspecting.
重点关注:All the operators are assigned to the DPU(see more details in 'inspect/inspect_DPUCZDX8G_ISA1_B4096.txt'),所有操作都分配给了DPU。
3. 其他有用函数
3.1 查看 torchvision 中模型
查看所有 torchvision.models 中的模型的方法:
import torchvision.models as models
dir(models)
---
['AlexNet',
'DenseNet',
'EfficientNet',
'GoogLeNet',
'GoogLeNetOutputs',
'Inception3',
'InceptionOutputs',
'MNASNet',
'MobileNetV2',
'MobileNetV3',
'RegNet',
'ResNet',
'ShuffleNetV2',
'SqueezeNet',
'VGG',
...
]
可以分类为以下模型:
AlexNet
DenseNet
EfficientNet
GoogLeNet
Inception3
MNASNet
MobileNetV2
MobileNetV3
RegNet
ResNet
ShuffleNetV2
SqueezeNet
VGG
每个模型名称后面可能还有具体的版本号或变种,如 resnet50 或 vgg16_bn,它们指的是同一模型的不同配置或使用批量归一化的版本。
3.2 保存模型
3.2.1 保存模型参数
torch.save(model.state_dict(), "my_model.pth")
- 这条语句只保存模型的参数(weights 和 biases),即模型的状态字典(state_dict)。
- state_dict 是一个 Python 字典,包含了模型中所有可学习参数的名称及其对应的张量值。
- 这种方法适用于只保存模型的参数,而不需要保存模型的结构时。
3.2.2 保存完整模型
torch.save(model, "my_model.pth")
- 这条语句保存整个模型对象,包括模型的结构和参数。
- 这种方法适用于需要保存完整的模型,包括其架构和状态,以便在加载时不需要重新定义模型结构。
检查方法:
print(isinstance(model, torch.nn.Module))
如果 model 继承自 torch.nn.Module,应该输出 True。
3.2.3 加载模型
model = torch.load(PATH)
4. 总结
本文详细介绍了如何使用 Vitis AI 的 Inspector 工具来检查和优化神经网络模型。首先,导入所需的库,包括 PyTorch 和 torchvision,并创建 Inspector 实例。接着,下载预训练的 ResNet18 模型,并生成一个虚拟输入张量。然后,使用 Inspector 对模型进行检查,指定计算设备和输出目录。检查过程包括解析模型、跟踪模型操作、权重均衡等步骤,最终生成检查报告和图像,显示所有操作已分配给 DPU。
此外,本文还介绍了如何查看 torchvision 中的所有模型,并分类列出常见的模型类型,如 AlexNet、DenseNet、ResNet 等。最后,讲解了保存和加载模型的方法,包括保存模型参数和完整模型,以及验证模型是否继承自 torch.nn.Module。
通过使用 Inspector 工具,用户可以在量化浮点模型前对其进行检查,生成的报告可用于指导模型优化,从而显著降低部署难度并缩短部署时间。这一过程对于在不同硬件架构上高效部署神经网络模型具有重要意义。