- cuda和cudnn都使用apt安装的版本，CUDA需要在bashrc设置环境变量
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)
# 创建模型实例
model = SimpleModel()
# 创建一个随机输入
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
