torch.onnx.export()

`torch.onnx.export()`是PyTorch提供的一个函数，用于将PyTorch模型转换为ONNX（开放神经网络交换格式）模型。这个功能允许你在保持模型结构不变的情况下将其部署到不支持PyTorch的环境中，比如移动设备、服务器或深度学习框架集成的服务。

以下是基本的调用语法：

from torch.onnx import export

def export_to_onnx(model, input_data, output_file, opset_version=11):
    # 确保模型是在评估模式下
    model.eval()
    
    # 设置为不跟踪梯度
    with torch.no_grad():
        traced_model = torch.jit.trace(model, input_data)  # 或者使用model.train().eval() if your model has dynamic shapes
        
    # 进行实际的ONNX导出
    torch.onnx.export(traced_model,  # 输入是TracedModule对象
                      input_data,  # 输入张量或数据集
                      output_file,  # 输出ONNX文件路径
                      opset_version=opset_version,  # 指定ONNX版本，保证兼容性
                      input_names=['input'],  # 输入名称列表
                      output_names=['output'],  # 输出名称列表
                      dynamic_axes={'input': {0: 'batch_size'},  # 动态轴设置，如batch维度
                                    'output': {0: 'batch_size'}})

export_to_onnx(model, input_data, 'model.onnx', opset_version=11)

其中，`opset_version`是一个关键参数，它表示ONNX规范的版本，选择较高的版本通常能得到更好的兼容性，但也可能导致更老版本的环境无法解析。

在使用`torch.onnx.export()`之前，确保你的PyTorch模型是静态图形式（常量形状），因为ONNX不支持动态计算图。如果模型的形状依赖于输入数据，你需要先通过`torch.jit.trace()`或者其他方式固定模型的形状。