- cuda和cudnn都使用apt安裝的版本憎账，CUDA需要在bashrc設置環(huán)境變量
export CUDA_HOME=/usr/local/cuda
export LD_LIBRARY_PATH=$CUDA_HOME/lib64:$LD_LIBRARY_PATH
export PATH=$CUDA_HOME/bin:$PATH
驅動不要升級，回退這個版本才能跑onnx
NVIDIA-Linux-x86_64-550.120.runonnx測試程序
import onnxruntime as ort
# 獲取可用的提供者
providers = ort.get_available_providers()
# 檢查是否支持 CUDA
if 'CUDAExecutionProvider' in providers:
print("onnxruntime 支持 CUDA")
print("可用的提供者:", providers)
else:
print("onnxruntime 不支持 CUDA")
print("可用的提供者:", providers)
# 導出簡單的onnx模型
import torch
import torch.nn as nn
# 定義一個簡單的模型
class SimpleModel(nn.Module):
def __init__(self):
super(SimpleModel, self).__init__()
self.fc = nn.Linear(10, 1)
def forward(self, x):
return self.fc(x)
# 創(chuàng)建模型實例
model = SimpleModel()
# 創(chuàng)建一個隨機輸入
x = torch.randn(1, 10)
# 導出為 ONNX 格式
torch.onnx.export(model, x, "dummy.onnx", input_names=["input"], output_names=["output"])
import onnxruntime as ort
# 獲取 CUDA 設備信息
providers = ort.get_available_providers()
if 'CUDAExecutionProvider' in providers:
sess_options = ort.SessionOptions()
sess = ort.InferenceSession("dummy.onnx", sess_options, providers=['CUDAExecutionProvider'])
cuda_devices = sess.get_providers()
print("可用的 CUDA 設備:", cuda_devices)
else:
print("onnxruntime 不支持 CUDA")
- labelme使用gpu加載onnx
https://blog.csdn.net/small_wu/article/details/137083509
