gh_mirrors/di/dino中的ONNX模型转换：从PyTorch到onnx格式的导出步骤-优快云博客

gh_mirrors/di/dino中的ONNX模型转换：从PyTorch到onnx格式的导出步骤

【免费下载链接】dino PyTorch code for Vision Transformers training with the Self-Supervised learning method DINO 项目地址: https://gitcode.com/gh_mirrors/di/dino

你还在为DINO模型部署时的框架兼容性问题烦恼吗？ONNX（Open Neural Network Exchange，开放神经网络交换格式）作为跨框架模型交互的桥梁，能帮你轻松解决这一痛点。本文将带你一步步完成从PyTorch模型到ONNX格式的转换，让你的DINO模型在各种部署环境中流畅运行。读完本文，你将掌握模型加载、参数配置、导出验证的全流程，并能应对常见问题。

准备工作

在开始转换前，确保你的环境满足以下要求：

PyTorch 1.8.0及以上版本（推荐使用项目兼容版本）
ONNX 1.9.0及以上版本
ONNX Runtime（可选，用于验证导出模型）

你可以通过项目的README.md获取详细的环境配置指南。如果需要安装ONNX相关依赖，可执行以下命令：

pip install onnx onnxruntime

模型结构解析

DINO项目的核心模型定义在vision_transformer.py中，主要包含VisionTransformer类及其衍生的模型变体。以常用的ViT-Small为例，模型通过vit_small函数初始化，包含以下关键参数：

patch_size：图像分块大小（默认16x16）
embed_dim：嵌入维度（384）
depth：Transformer块数量（12）
num_heads：注意力头数量（6）

# VisionTransformer类核心结构（源自vision_transformer.py）
class VisionTransformer(nn.Module):
    def __init__(self, img_size=[224], patch_size=16, in_chans=3, num_classes=0, 
                 embed_dim=768, depth=12, num_heads=12, mlp_ratio=4.):
        super().__init__()
        self.patch_embed = PatchEmbed(img_size[0], patch_size, in_chans, embed_dim)
        self.cls_token = nn.Parameter(torch.zeros(1, 1, embed_dim))
        self.pos_embed = nn.Parameter(torch.zeros(1, num_patches + 1, embed_dim))
        self.blocks = nn.ModuleList([Block(dim=embed_dim, num_heads=num_heads) for _ in range(depth)])
        self.norm = nn.LayerNorm(embed_dim)
        self.head = nn.Linear(embed_dim, num_classes) if num_classes > 0 else nn.Identity()

模型加载可通过hubconf.py中的预定义函数实现，例如加载ViT-Small/16模型：

model = torch.hub.load('facebookresearch/dino:main', 'dino_vits16')

转换步骤

步骤1：加载预训练模型

使用项目提供的模型加载工具加载预训练权重。以下代码示例展示如何加载ViT-Small/16模型：

import torch
import vision_transformer as vits

# 加载模型架构
model = vits.vit_small(patch_size=16, num_classes=0)
# 加载预训练权重（可从项目指定地址下载）
state_dict = torch.load("dino_deitsmall16_pretrain.pth", map_location="cpu")
model.load_state_dict(state_dict, strict=True)
model.eval()  # 设置为评估模式

步骤2：准备输入张量

根据模型要求创建输入张量，ViT-Small/16默认输入尺寸为224x224x3：

# 创建随机输入张量（批量大小=1，通道=3，高度=224，宽度=224）
input_tensor = torch.randn(1, 3, 224, 224)

步骤3：导出ONNX模型

使用torch.onnx.export函数将PyTorch模型导出为ONNX格式。关键参数说明：

input_names：输入节点名称
output_names：输出节点名称
dynamic_axes：支持动态批次大小

# 导出模型
torch.onnx.export(
    model,                      # 模型实例
    input_tensor,               # 输入张量
    "dino_vits16.onnx",         # 输出文件路径
    input_names=["input"],      # 输入节点名称
    output_names=["output"],    # 输出节点名称
    dynamic_axes={              # 动态维度设置
        "input": {0: "batch_size"},
        "output": {0: "batch_size"}
    },
    opset_version=12            # ONNX算子集版本
)

步骤4：验证导出模型

使用ONNX Runtime加载导出的模型并验证输出一致性：

import onnxruntime as ort
import numpy as np

# 加载ONNX模型
ort_session = ort.InferenceSession("dino_vits16.onnx")
# 准备输入数据（转换为numpy数组）
input_np = input_tensor.numpy()
# 推理
ort_output = ort_session.run(None, {"input": input_np})[0]
# PyTorch推理
with torch.no_grad():
    torch_output = model(input_tensor).numpy()
# 验证输出差异
np.testing.assert_allclose(torch_output, ort_output, rtol=1e-3, atol=1e-3)
print("模型验证成功，输出差异在可接受范围内")

常见问题与解决方案

问题1：导出时出现不支持的算子

原因：某些PyTorch算子在ONNX中没有直接对应实现。
解决方案：更新PyTorch版本或使用torch.onnx.export的opset_version参数指定更高版本（如12+）。

问题2：动态输入尺寸导致推理失败

原因：模型中包含固定形状的操作（如位置编码）。
解决方案：在vision_transformer.py的interpolate_pos_encoding方法中确保位置编码支持动态插值，或导出时固定输入尺寸。

问题3：导出模型体积过大

解决方案：使用ONNX优化工具（如onnxsim）简化模型：

pip install onnx-simplifier
python -m onnxsim dino_vits16.onnx dino_vits16_simplified.onnx

总结与展望

通过本文的步骤，你已成功将DINO的PyTorch模型转换为ONNX格式。ONNX模型可进一步部署到各种平台，如TensorRT、OpenVINO等，实现高效推理。项目后续可能会集成更便捷的导出工具，你可以关注main_dino.py中的更新或参与贡献。

如果你在转换过程中遇到其他问题，欢迎查阅项目文档或提交issue。祝你部署顺利！

gh_mirrors/di/dino中的ONNX模型转换：从PyTorch到onnx格式的导出步骤