PyTorch模型部署二、动态维度解释
·
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导出成功!")
更多推荐
所有评论(0)