跳至主要內容

DeepSeek AI 模型性能测试脚本

Gody2025年4月3日大约 8 分钟

注

开源声明: 本脚本由 Gody 原创开发，采用 MIT 许可证开源分享。

Copyright (c) 2025 Gody

使用、修改和分发本代码时请保留此声明。

DeepSeek AI 模型性能测试脚本

这个脚本用于对比测试 DeepSeek AI 模型在 GPU 和 CPU 模式下的性能差异。通过 Docker 运行 Ollama 容器，分别在 GPU 和 CPU 环境下执行相同的查询，并比较两种模式下的处理速度。

演示视频

A BiliBili video

点击上方视频查看 DeepSeek AI 模型的性能测试演示

功能特点

  • 自动检测 GPU 是否可用
  • 在 GPU 和 CPU 模式下分别测试 AI 模型性能
  • 比较两种模式的 token 处理速度
  • 生成详细的测试报告日志

使用方法

./model_test.sh <模型名称> <问题内容>

示例:

./model_test.sh deepseek-r1:8b "请介绍一下量子计算"

脚本代码

#!/bin/bash

# AI模型性能测试脚本（GPU vs 内存模式）
# 用法: ./model_test.sh <模型名称> <问题内容>

# 可设置的内存变量（默认值：60G）
MEMORY_LIMIT="60g"

# 颜色设置（仅用于终端显示）
RED="\033[0;31m"
GREEN="\033[0;32m"
YELLOW="\033[0;33m"
BLUE="\033[0;34m"
NC="\033[0m" # 无颜色

# 检查参数
if [ $# -lt 2 ]; then
    echo -e "${RED}错误：需要提供模型名称和问题内容${NC}"
    echo -e "用法: ${GREEN}./model_test.sh <模型名称> <问题内容>${NC}"
    echo -e "示例: ${GREEN}./model_test.sh deepseek-r1:8b \"请介绍一下量子计算\"${NC}"
    exit 1
fi

# 获取参数
MODEL="$1"
QUESTION="$2"

# 创建日志文件 - 先创建文件再写入内容
TIMESTAMP=$(date +"%Y%m%d_%H%M%S")
LOG_FILE="${MODEL//[:\/]/_}_benchmark_${TIMESTAMP}.log"

# 添加测试信息到日志文件开头
{
    echo "===== AI模型性能测试报告 ====="
    echo "测试时间: $(date)"
    echo "模型: $MODEL"
    echo "问题: $QUESTION"
    echo "==========================="
    echo ""
} > "$LOG_FILE"

# 日志记录（无颜色）
log_to_file() {
    echo "$1" >> "$LOG_FILE"
}

# log_message函数 - 优化版本
log_message() {
    # 终端显示（带颜色）
    echo -e "$1"
    
    # 直接使用sed替换掉所有颜色代码，简化为一步操作
    sed 's/\x1B\[[0-9;]*[mK]//g; s/\\033\[[0-9;]*[mK]//g' <<< "$1" >> "$LOG_FILE"
}

# 函数：清理Docker容器 - 优化版本
cleanup_docker() {
    log_message "${YELLOW}清理Ollama容器...${NC}"
    
    # 检查容器是否存在并运行
    if docker ps -q --filter "name=ollama" | grep -q .; then
        log_message "${YELLOW}停止现有容器...${NC}"
        docker stop ollama >/dev/null 2>&1
        docker rm -f ollama >/dev/null 2>&1 || true
    elif docker ps -a -q --filter "name=ollama" | grep -q .; then
        log_message "${YELLOW}清理未运行的容器...${NC}"
        docker rm -f ollama >/dev/null 2>&1 || true
    else
        log_message "${YELLOW}没有发现现有的Ollama容器${NC}"
    fi
    
    # 确保挂载卷存在
    if ! docker volume inspect ollama >/dev/null 2>&1; then
        log_message "${YELLOW}创建ollama数据卷...${NC}"
        docker volume create ollama >/dev/null 2>&1
    fi
    
    sleep 2
}

# 检查GPU是否可用 - 优化版本
check_gpu() {
    log_message "${YELLOW}检查GPU是否可用...${NC}"
    
    # 缓存nvidia-smi命令结果，避免重复执行
    if ! command -v nvidia-smi &> /dev/null || ! nvidia-smi &> /dev/null; then
        log_message "${RED}警告：GPU可能不可用${NC}"
        return 1
    fi
    
    log_message "${GREEN}GPU检测正常${NC}"
    return 0
}

# 函数：等待API可用 - 优化版本
wait_for_api() {
    log_message "${YELLOW}等待Ollama API可用...${NC}"
    
    # 使用printf代替多次echo，减少进程创建
    for i in {1..30}; do
        if curl -s http://localhost:11434/api/version >/dev/null 2>&1; then
            log_message "${GREEN}API已就绪${NC}"
            return 0
        fi
        sleep 1
        printf "."
    done
    printf "\n"
    
    log_message "${RED}错误：API在30秒内未响应${NC}"
    return 1
}

# 函数：检查模型是否已下载
check_model() {
    log_message "${YELLOW}检查模型 $MODEL 是否已下载...${NC}"
    
    # 缓存模型列表，避免重复查询
    local MODEL_LIST
    MODEL_LIST=$(docker exec ollama ollama list 2>/dev/null)
    
    if ! echo "$MODEL_LIST" | grep -q "$MODEL"; then
        log_message "${YELLOW}下载模型 $MODEL 中...${NC}"
        docker exec ollama ollama pull "$MODEL"
        
        if [ $? -ne 0 ]; then
            log_message "${RED}错误：无法下载模型 $MODEL${NC}"
            return 1
        fi
        log_message "${GREEN}模型 $MODEL 下载完成${NC}"
    else
        log_message "${GREEN}模型 $MODEL 已存在${NC}"
    fi
    return 0
}

# 优化版运行测试函数
run_test() {
    local mode=$1
    
    log_message "\n${BLUE}===== $mode模式测试 =====${NC}"
    log_message "${YELLOW}开始测试，请稍候...${NC}"
    
    # 临时文件用于捕获性能数据
    TEMP_FILE=$(mktemp)
    
    # 使用echo和管道传递问题
    log_message "${YELLOW}运行模型中...${NC}"
    echo "$QUESTION" | docker exec -i ollama ollama run $MODEL --verbose 2> "$TEMP_FILE"
    
    # 使用单次读取文件提取所有性能数据，减少重复IO操作
    local PERF_DATA
    PERF_DATA=$(cat "$TEMP_FILE")
    
    # 提取性能数据 - 使用单次处理代替多次grep
    TOKEN_SPEED=$(echo "$PERF_DATA" | grep -o "[0-9.]\+ tokens/s" | tail -1)
    TOTAL_TOKENS=$(echo "$PERF_DATA" | grep -o "total [0-9]\+ tokens" | tail -1)
    ELAPSED_TIME=$(echo "$PERF_DATA" | grep -o "in [0-9.]\+s" | tail -1)
    
    # 记录性能数据
    if [ -n "$TOKEN_SPEED" ]; then
        log_message "${GREEN}$mode模式token速度: $TOKEN_SPEED${NC}"
    else
        log_message "${RED}未能获取token速度数据${NC}"
        log_to_file "未能获取token速度数据，错误输出:"
        # 使用head和管道，避免多次文件读取
        head -10 "$TEMP_FILE" | while read -r line; do
            log_to_file "$line"
        done
    fi
    
    if [ -n "$TOTAL_TOKENS" ]; then
        log_message "${GREEN}$TOTAL_TOKENS${NC}"
    fi
    
    if [ -n "$ELAPSED_TIME" ]; then
        log_message "${GREEN}用时: $ELAPSED_TIME${NC}"
    fi
    
    # 清理临时文件
    rm -f "$TEMP_FILE"
}

# 主测试流程
log_message "${BLUE}开始AI模型性能测试...${NC}"

# 预先准备模型，避免重复下载
log_message "\n${YELLOW}预先检查模型是否已下载...${NC}"
cleanup_docker

# 启动临时容器来预下载模型
log_message "${YELLOW}启动临时容器准备模型...${NC}"
docker run -d -v ollama:/root/.ollama -p 11434:11434 --name ollama ollama/ollama serve

# 等待API可用并确保模型已下载
if wait_for_api; then
    if check_model; then
        log_message "${GREEN}模型已准备就绪，将继续进行测试${NC}"
    else
        log_message "${RED}模型准备失败，但会继续尝试测试${NC}"
    fi
else
    log_message "${RED}API准备失败，但会继续尝试测试${NC}"
fi

# 停止临时容器
log_message "${YELLOW}清理临时容器...${NC}"
docker stop ollama >/dev/null 2>&1
docker rm -f ollama >/dev/null 2>&1

# ===== GPU测试部分 =====
log_message "\n${GREEN}========== 启动GPU模式测试 ==========${NC}"

# 检查GPU可用性
if ! check_gpu; then
    log_message "${YELLOW}GPU可能不可用，但仍将尝试GPU模式测试${NC}"
fi

# 清理并启动GPU容器
cleanup_docker
log_message "${YELLOW}启动GPU模式的Ollama容器...${NC}"
docker run -d --gpus=all -v ollama:/root/.ollama -p 11434:11434 --name ollama ollama/ollama serve

# 等待API和检查模型
wait_for_api && check_model && run_test "GPU" || log_message "${RED}GPU测试失败${NC}"

# ===== CPU测试部分 =====
log_message "\n${GREEN}========== 启动CPU模式测试 ==========${NC}"

# 清理并启动CPU容器（纯内存模式）
cleanup_docker
log_message "${YELLOW}启动CPU模式的Ollama容器...${NC}"
docker run -d -v ollama:/root/.ollama -p 11434:11434 --memory=${MEMORY_LIMIT} --cpus=16 --name ollama ollama/ollama serve

# 等待API和检查模型
wait_for_api && check_model && run_test "CPU" || log_message "${RED}CPU测试失败${NC}"

# 清理容器
cleanup_docker

# 总结结果 - 优化数据提取方式
log_message "\n${BLUE}========== 性能对比结果 ==========${NC}"
log_to_file "\n========== 性能对比结果 =========="

# 使用更精确的grep模式提取速度值
GPU_SPEED=$(grep -o "GPU模式token速度: [0-9.]\+ tokens/s" "$LOG_FILE" | grep -o "[0-9.]\+" | head -1)
CPU_SPEED=$(grep -o "CPU模式token速度: [0-9.]\+ tokens/s" "$LOG_FILE" | grep -o "[0-9.]\+" | head -1)

if [ -n "$GPU_SPEED" ] && [ -n "$CPU_SPEED" ] && [ "$GPU_SPEED" != "0" ] && [ "$CPU_SPEED" != "0" ]; then
    # 确保bc已安装，缓存检查结果
    if ! command -v bc &> /dev/null; then
        log_message "${YELLOW}安装bc计算工具...${NC}"
        apt-get update && apt-get install -y bc
    fi
    
    SPEEDUP=$(echo "scale=2; $GPU_SPEED / $CPU_SPEED" | bc)
    
    # 使用数组和循环减少重复代码
    for msg in "GPU速度: $GPU_SPEED tokens/s" "CPU速度: $CPU_SPEED tokens/s" "GPU比CPU快约 $SPEEDUP 倍"; do
        log_message "${GREEN}$msg${NC}"
        log_to_file "$msg"
    done
else
    log_message "${RED}无法计算性能比较，数据不完整${NC}"
    log_to_file "无法计算性能比较，数据不完整"
fi

log_message "\n${GREEN}测试完成！结果已保存到: $LOG_FILE${NC}"

环境要求

必需组件

Docker

  • 版本要求: 20.10.0 或更高
  • 配置要求: 需要正确安装和配置 Docker 引擎
  • 验证方法: docker --version 命令应成功返回版本信息
  • 权限要求: 运行脚本的用户需要有 docker 命令的执行权限

提示

如果不是 docker 组成员，可使用 sudo usermod -aG docker $USER 添加当前用户，然后重新登录

NVIDIA GPU 硬件与驱动 (用于 GPU 测试)

  • 硬件要求: NVIDIA GPU 显卡，最低 4GB 显存，推荐 8GB+
  • 驱动要求: NVIDIA 驱动 450.80.02 或更高版本
  • 软件要求: CUDA 11.4+ 和 NVIDIA Container Toolkit
  • 验证方法: nvidia-smi 命令能正确显示 GPU 信息

注

无 GPU 环境下仍可运行，但将只执行 CPU 测试部分

命令行工具

  • bc: 用于计算性能比率（脚本会自动安装）
  • curl: 用于 API 测试（大多数系统已预装）

系统要求

硬件配置

  • 处理器: 8+ 核心 CPU (用于 CPU 模式测试)
  • 内存:
    • 最低: 16GB RAM
    • 推荐: 32GB+ RAM (脚本默认限制为 60GB)
  • 存储:
    • 可用空间: 20GB+ (用于 Docker 镜像和模型)
    • 类型: SSD 可显著提升性能

网络要求

  • 连接: 稳定的互联网连接
  • 带宽: 建议 10Mbps+ (模型下载需要数 GB 流量)
  • 防火墙: 允许 Docker 容器网络访问

操作系统兼容性

  • 推荐系统: Ubuntu 20.04/22.04 LTS
  • 兼容系统: Debian、CentOS、Fedora (可能需要调整包管理命令)
  • 有限支持: macOS/Windows WSL2 (GPU 直通受限)

注意

在非 Linux 环境下，GPU 功能可能无法正常工作

注意事项

重要

首次运行将下载大型模型文件，可能需要较长时间

  • 对于大型模型 (34B 版本)，建议使用 16GB+ VRAM 的 GPU
  • 可通过修改脚本中的 MEMORY_LIMIT 变量调整内存限制
  • 确保 Docker 具有足够的系统资源配额和权限
  • 测试结果会保存到与模型同名的日志文件中

贡献者

gody