mirror of
https://github.com/deepinsight/insightface.git
synced 2025-12-30 08:02:27 +00:00
66 lines
2.4 KiB
Bash
66 lines
2.4 KiB
Bash
#!/bin/bash
|
||
|
||
# 函数：获取CUDA和Ubuntu版本标签
|
||
# 如果CUDA_TAG环境变量已设置，则使用该值
|
||
# 否则自动检测CUDA和Ubuntu版本并生成标签
|
||
# 格式: cudaXX_ubuntuXX.XX
|
||
# 如果检测不到某个版本，则用"none"代替
|
||
get_cuda_ubuntu_tag() {
|
||
# 如果CUDA_TAG已设置，则直接返回
|
||
if [ -n "${CUDA_TAG}" ]; then
|
||
echo "${CUDA_TAG}"
|
||
return 0
|
||
fi
|
||
|
||
# 获取CUDA版本
|
||
CUDA_VERSION="_none"
|
||
if command -v nvcc &> /dev/null; then
|
||
# 尝试从nvcc获取版本
|
||
CUDA_VERSION=$(nvcc --version 2>/dev/null | grep "release" | awk '{print $6}' | cut -d',' -f1 | tr -d '.')
|
||
if [ -z "${CUDA_VERSION}" ]; then
|
||
CUDA_VERSION="_none"
|
||
else
|
||
CUDA_VERSION="${CUDA_VERSION}"
|
||
fi
|
||
elif [ -f "/usr/local/cuda/version.txt" ]; then
|
||
# 尝试从CUDA安装目录获取版本
|
||
CUDA_VERSION=$(cat /usr/local/cuda/version.txt 2>/dev/null | grep "CUDA Version" | awk '{print $3}' | tr -d '.')
|
||
if [ -z "${CUDA_VERSION}" ]; then
|
||
CUDA_VERSION="_none"
|
||
fi
|
||
elif [ -d "/usr/local/cuda" ] && ls -l /usr/local/cuda 2>/dev/null | grep -q "cuda-"; then
|
||
# 尝试从符号链接获取版本
|
||
CUDA_LINK=$(ls -l /usr/local/cuda 2>/dev/null | grep -o "cuda-[0-9.]*" | head -n 1)
|
||
CUDA_VERSION=$(echo "${CUDA_LINK}" | cut -d'-' -f2 | tr -d '.')
|
||
if [ -z "${CUDA_VERSION}" ]; then
|
||
CUDA_VERSION="_none"
|
||
fi
|
||
fi
|
||
|
||
# 获取Ubuntu版本
|
||
UBUNTU_VERSION="_none"
|
||
if [ -f "/etc/os-release" ]; then
|
||
# 检查是否是Ubuntu
|
||
if grep -q "Ubuntu" /etc/os-release 2>/dev/null; then
|
||
UBUNTU_VERSION=$(grep "VERSION_ID" /etc/os-release 2>/dev/null | cut -d'"' -f2)
|
||
if [ -z "${UBUNTU_VERSION}" ]; then
|
||
UBUNTU_VERSION="_none"
|
||
fi
|
||
fi
|
||
elif [ -f "/etc/lsb-release" ]; then
|
||
# 尝试从lsb-release获取版本
|
||
if grep -q "Ubuntu" /etc/lsb-release 2>/dev/null; then
|
||
UBUNTU_VERSION=$(grep "DISTRIB_RELEASE" /etc/lsb-release 2>/dev/null | cut -d'=' -f2)
|
||
if [ -z "${UBUNTU_VERSION}" ]; then
|
||
UBUNTU_VERSION="_none"
|
||
fi
|
||
fi
|
||
fi
|
||
|
||
# 生成并返回标签
|
||
echo "cuda${CUDA_VERSION}_ubuntu${UBUNTU_VERSION}"
|
||
}
|
||
|
||
# 使用示例
|
||
CUDA_TAG=$(get_cuda_ubuntu_tag)
|
||
echo "Generated tag: ${CUDA_TAG}" |