Pytorch 转 ONNX

pip install onnx onnxruntime onnxruntime-gpu

====
PyTorch model to ONNX:
    model = models.efficientnet_v2_s()
    model.to(device)

    torch_input = torch.randn(1, 3, 384, 384, device="cuda")
    torch.onnx.export(model,  # model being run
                      torch_input,  # model input
                      "./effnet.onnx",
                      # where to save the model (can be a file or file-like object)
                      opset_version=11,  # the ONNX version to export the model to
                      input_names=['input'],  # the model's input names
                      output_names=['output']  # the model's output names
                      )

====
Test ONNX model:
import multiprocessing
import onnx
import onnxruntime as ort
import numpy as np

def main():
    onnx_model_path = "effnet.onnx"

    # Load the ONNX model
    model = onnx.load(onnx_model_path)

    # Check that the model is well formed
    onnx.checker.check_model(model)

    # Print a human readable representation of the graph
    #print(onnx.helper.printable_graph(model.graph))

    #onnx_provider = 'CPUExecutionProvider'
    onnx_provider = 'CUDAExecutionProvider'
    ort_session = ort.InferenceSession(onnx_model_path, providers=[onnx_provider])

    outputs = ort_session.run(
        None,
        {"input": np.random.randn(1, 3, 384, 384).astype(np.float32)},
    )
    #print(outputs)
    #print(outputs[0].shape)
    #print(outputs[0])
    print(np.argmax(outputs[0]))

if __name__ == "__main__":
    multiprocessing.freeze_support()
    main()


====
Error:
EP Error C:\a\_work\1\s\onnxruntime\python\onnxruntime_pybind_state.cc:866 onnxruntime::python::CreateExecutionProviderInstance CUDA_PATH is set but CUDA wasnt able to be loaded. Please install the correct version of CUDA andcuDNN as mentioned in the GPU requirements page  (https://onnxruntime.ai/docs/execution-providers/CUDA-ExecutionProvider.html#requirements),  make sure they're in the PATH, and that your GPU is supported.
 when using ['CUDAExecutionProvider']

https://onnxruntime.ai/docs/execution-providers/CUDA-ExecutionProvider.html#requirements
ONNX Runtime	CUDA	cuDNN	                              Notes
1.18	        12.4	8.9.2.26 (Linux)/8.9.2.26 (Windows)   The default CUDA version for ORT 1.18 is CUDA 11.8. 
                                                              To install CUDA 12 package, please look at Install ORT. Java CUDA 12 support is back for release 1.18
1.18	        11.8	8.9.2.26 (Linux)/8.9.2.26 (Windows)

需要安装CUDA 12.4 or 11.8 + cuDNN 8.9.2.26：
https://developer.nvidia.com/cuda-12-4-1-download-archive?target_os=Windows&target_arch=x86_64&target_version=10&target_type=exe_local
https://developer.nvidia.com/rdp/cudnn-archive

Error:
Could not locate zlibwapi.dll. Please make sure it is in your library path!
https://docs.nvidia.com/deeplearning/cudnn/archives/cudnn-892/index.html
http://www.winimage.com/zLibDll/zlib123dllx64.zip