def export_onnx(model, im, file, opset=11, dynamic=False, simplify=False, prefix="ONNX:"):
    """
    导出 EfficientNetV2 模型为 ONNX 格式，支持动态轴和简化选项。
    """
    print(f"\n{prefix} Starting export with ONNX...")

    f = str(file.with_suffix(".onnx"))

    dynamic_axes = {"input": {0: "batch_size", 2: "height", 3: "width"},
                    "output": {0: "batch_size"}} if dynamic else None

    torch.onnx.export(
        model,
        im,
        f,
        opset_version=opset,
        do_constant_folding=True,
        input_names=["input"],
        output_names=["output"],
        dynamic_axes=dynamic_axes,
        verbose=False
    )

    model_onnx = onnx.load(f)
    onnx.checker.check_model(model_onnx)
    print(f"{prefix} Model exported to {f}")

    if simplify:
        try:
            from onnxsim import simplify
            model_onnx, check = simplify(model_onnx)
            assert check, "Simplified ONNX model could not be validated"
            onnx.save(model_onnx, f)
            print(f"{prefix} Model simplified and saved to {f}")
        except Exception as e:
            print(f"{prefix} Simplification failed: {e}")

    return f, model_onnx

dummy_input = torch.randn(1, 3, 224, 224).to(device)
# 导出 ONNX 模型
export_onnx(model, dummy_input, file=file_path, opset=11, dynamic=True, simplify=True)