DenseNet ONNX (从 TorchVision 导出)

本仓库托管了 TorchVision DenseNet ImageNet 分类器的 ONNX 导出版本。模型来源于 Torch Hub (pytorch/vision:v0.10.0) 并导出为 ONNX 格式,以便使用 ONNX Runtime 轻松部署。

• 源权重:TorchVision v0.10.0 (ImageNet-1k)
• 支持的变体:densenet121、densenet161、densenet169
• 默认输入:NCHW 1 x 3 x 224 x 224 RGB,使用 ImageNet 标准化

如果您在 Hugging Face Hub 上浏览,可以在 Files 标签页下找到 ONNX 文件。您也可以使用下面的脚本自行重现导出过程。

快速开始:使用 ONNX Runtime 推理

import numpy as np
import onnxruntime as ort
import torch
from PIL import Image
from torchvision.transforms import Compose, Resize, CenterCrop, PILToTensor, ConvertImageDtype, Normalize

# 1) 加载 ONNX 模型
session = ort.InferenceSession("densenet121.onnx", providers=["CPUExecutionProvider"])
input_name = session.get_inputs()[0].name
output_name = session.get_outputs()[0].name

# 2) 预处理 (ImageNet)
img = Image.open("your_image.jpg").convert("RGB")
transform = Compose([
    Resize(256),
    CenterCrop(224),
    PILToTensor(),
    ConvertImageDtype(torch.float32),
    Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
])
inp = transform(img).unsqueeze(0).numpy()  # NCHW float32

# 3) 运行推理
scores = session.run([output_name], {input_name: inp})[0]  # (1, 1000)
pred = int(scores.argmax(axis=1)[0])
print("Top-1 class index:", pred)

要将索引映射到 ImageNet-1k 标签,您可以获取: https://raw.githubusercontent.com/pytorch/hub/master/imagenet_classes.txt

从 Torch Hub 导出到 ONNX (脚本)

使用提供的导出脚本从 Torch Hub 重现 ONNX 文件:

python export_to_onnx.py --variant densenet121 --output densenet121.onnx --opset 17

选项:

• --variant:densenet121、densenet161、densenet169 之一
• --opset:ONNX opset 版本(默认:17)
• --static-batch:使批次大小为静态(默认情况下,导出使用动态批次)
• --no-sample:使用虚拟输入而非样本图像进行导出

最小化导出

import torch
from PIL import Image
from torchvision.transforms import Compose, Resize, CenterCrop, PILToTensor, ConvertImageDtype, Normalize

variant = "densenet121"  # 或 161/169
onnx_path = "densenet121.onnx"

torch_model = torch.hub.load("pytorch/vision:v0.10.0", variant, pretrained=True).eval()

img = Image.open("your_image.jpg").convert("RGB")
transform = Compose([
    Resize(256),
    CenterCrop(224),
    PILToTensor(),
    ConvertImageDtype(torch.float32),
    Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
])
img_tensor = transform(img).unsqueeze(0).contiguous()

torch.onnx.export(
    torch_model,
    img_tensor,
    onnx_path,
    opset_version=17,
    input_names=["input"],
    output_names=["logits"],
    dynamic_axes={"input": {0: "batch_size"}, "logits": {0: "batch_size"}},
)