PyTorch模型部署一、模型转换为onnx
https://blog.csdn.net/qq_41008520/article/details/159473715?spm=1001.2014.3001.5501

1. 警告信息解释

UserWarning: ‘dynamic_axes’ is not recommended when dynamo=True

•新版本的PyTorch(>=2.0)使用了新的ONNX导出引擎(称为"dynamo")

• 为什么是警告而不是错误:系统仍然支持dynamic_axes参数,但推荐使用新的dynamic_shapes参数

2. 动态维度(dynamic_axes)详解

什么是动态维度?

动态维度允许模型接受不同尺寸的输入。在你的例子中:
dynamic_axes={
‘input’: {0: ‘batch_size’}, # 第0维(batch维度)是动态的
‘output’: {0: ‘batch_size’} # 输出的batch维度也是动态的
}

这意味着:
• 可以输入 torch.randn(1, 10) → 输出形状 (1, 5)

• 可以输入 torch.randn(4, 10) → 输出形状 (4, 5)

• 可以输入 torch.randn(16, 10) → 输出形状 (16, 5)

batch维度动态,但特征维度固定为10。

静态 vs 动态维度对比

静态维度(无dynamic_axes):

# 导出时只支持batch_size=1
torch.onnx.export(
    ...,
    # 不设置dynamic_axes
)
# 只能推理batch_size=1的输入

动态维度(有dynamic_axes):

# 导出时支持任意batch_size
torch.onnx.export(
    ...,
    dynamic_axes={
        'input': {0: 'batch_size'},  # 第0维可变
        'output': {0: 'batch_size'}
    }
)
# 可以推理batch_size=1, 2, 4, 8, 16...的输入

验证导出的动态维度模型

创建验证脚本来确认模型正常工作:

# verify_onnx.py
import torch
import torch.nn as nn
import onnx
import onnxruntime as ort
import numpy as np

# 加载ONNX模型
onnx_model_path = "simple_model.onnx"

# 1. 检查模型格式
print("1. 检查ONNX模型格式...")
onnx_model = onnx.load(onnx_model_path)
onnx.checker.check_model(onnx_model)
print("✓ ONNX模型格式正确")

# 2. 查看模型信息
print(f"\n2. 模型信息:")
print(f"   IR版本: {onnx_model.ir_version}")
print(f"   生产者: {onnx_model.producer_name}")
print(f"   Opset版本: {onnx_model.opset_import[0].version}")

# 查看输入输出形状
print(f"\n3. 输入输出信息:")
for input in onnx_model.graph.input:
    print(f"   输入: {input.name}")
    for dim in input.type.tensor_type.shape.dim:
        if dim.dim_param:  # 动态维度
            print(f"      - 维度: {dim.dim_param} (动态)")
        else:  # 固定维度
            print(f"      - 维度: {dim.dim_value} (固定)")

for output in onnx_model.graph.output:
    print(f"   输出: {output.name}")
    for dim in output.type.tensor_type.shape.dim:
        if dim.dim_param:
            print(f"      - 维度: {dim.dim_param} (动态)")
        else:
            print(f"      - 维度: {dim.dim_value} (固定)")

# 3. 创建原始PyTorch模型(用于对比)
class SimpleModel(nn.Module):
    def __init__(self):
        super(SimpleModel, self).__init__()
        self.linear = nn.Linear(10, 5)
        self.relu = nn.ReLU()
    
    def forward(self, x):
        return self.relu(self.linear(x))

torch_model = SimpleModel()
torch_model.eval()

# 4. 测试不同batch_size的推理
print("\n4. 测试不同batch_size的推理:")
ort_session = ort.InferenceSession(onnx_model_path)
input_name = ort_session.get_inputs()[0].name
output_name = ort_session.get_outputs()[0].name

batch_sizes = [1, 2, 4, 8]

for batch_size in batch_sizes:
    print(f"\n  Batch size = {batch_size}:")
    
    # 创建输入
    torch_input = torch.randn(batch_size, 10)
    numpy_input = torch_input.numpy()
    
    # PyTorch推理
    with torch.no_grad():
        torch_output = torch_model(torch_input).numpy()
    
    # ONNX推理
    ort_output = ort_session.run([output_name], {input_name: numpy_input})[0]
    
    # 对比结果
    diff = np.abs(torch_output - ort_output).max()
    print(f"    PyTorch输出形状: {torch_output.shape}")
    print(f"    ONNX输出形状: {ort_output.shape}")
    print(f"    最大差异: {diff:.8f}")
    
    if diff < 1e-5:
        print("    ✓ 结果一致")
    else:
        print(f"    ⚠ 有差异,但可能在误差范围内")

print("\n" + "="*50)
print("验证完成!模型可以正常工作。")
print("动态batch_size测试通过!")
print("="*50)

运行验证:
python3 verify_onnx.py

运行这个脚本来确认一切正常:

# final_check.py
import onnxruntime as ort
import numpy as np

# 测试导出的模型
print("测试导出的simple_model.onnx...")
session = ort.InferenceSession("simple_model.onnx")

# 测试不同batch_size
test_cases = [
    (1, "单样本推理"),
    (4, "小批量推理"),
    (16, "大批量推理")
]

for batch_size, description in test_cases:
    # 准备输入
    input_data = np.random.randn(batch_size, 10).astype(np.float32)
    
    # 推理
    outputs = session.run(None, {'input': input_data})
    
    print(f"\n{description} (batch_size={batch_size}):")
    print(f"  输入形状: {input_data.shape}")
    print(f"  输出形状: {outputs[0].shape}")
    print(f"  输出范围: [{outputs[0].min():.4f}, {outputs[0].max():.4f}]")

print("\n✅ 模型可以正常处理不同batch_size的输入!")
print("✅ ONNX导出成功!")
Logo

腾讯云面向开发者汇聚海量精品云计算使用和开发经验,营造开放的云计算技术生态圈。

更多推荐