AI编程 性能优化教程｜附完整命令

在 AI 编程项目中，很多开发者一开始关注的是“模型能不能跑起来”，但当项目进入真实业务场景后，性能问题会迅速暴露出来：接口响应慢、GPU 利用率低、显存频繁爆掉、推理成本过高、并发一上来服务就崩、训练时间过长、数据加载成为瓶颈等。

性能优化不是单点技巧，而是一套系统工程。本文将从 环境准备、性能分析、数据加载优化、模型训练优化、推理加速、服务部署、并发压测、监控与排查 等方面，系统讲解 AI 编程中的性能优化方法，并附上可直接使用的完整命令。

一、AI 项目性能优化的核心目标

在开始优化之前，需要先明确优化目标。不同阶段的 AI 项目，性能指标并不一样。

1. 训练阶段关注什么？

训练阶段通常关注以下指标：

• 训练速度：每秒处理多少样本，或者每个 epoch 耗时多久；
• GPU 利用率：GPU 是否真正被充分使用；
• 显存占用：模型、batch size、优化器状态是否导致显存溢出；
• 数据加载速度：CPU、磁盘、网络是否拖慢训练；
• 分布式效率：多卡训练是否接近线性加速；
• 收敛效果：不能为了速度牺牲模型精度。

2. 推理阶段关注什么？

推理阶段更关注：

• 延迟 Latency：单次请求从进入到返回需要多久；
• 吞吐 Throughput：单位时间内能处理多少请求；
• 并发能力：同时有多个请求时服务是否稳定；
• 资源成本：每千次请求消耗多少 GPU/CPU 成本；
• 冷启动时间：模型加载和服务启动是否太慢；
• 稳定性：长时间运行是否内存泄漏、显存泄漏。

二、环境准备：先把基础工具装好

优化性能的第一步不是盲目改代码，而是先学会观察。你需要安装一些常用工具，用来查看 CPU、GPU、内存、磁盘和网络情况。

1. 查看系统信息

uname -a
lsb_release -a
lscpu
free -h
df -h

2. 查看 GPU 信息

如果使用 NVIDIA GPU，可以执行：

nvidia-smi

持续监控 GPU：

watch -n 1 nvidia-smi

查看更详细的 GPU 进程信息：

nvidia-smi pmon -c 1

查看显存、功耗、温度等：

nvidia-smi --query-gpu=name,memory.total,memory.used,utilization.gpu,temperature.gpu,power.draw --format=csv

3. 安装常用性能工具

Ubuntu/Debian 系统：

sudo apt update
sudo apt install -y htop iotop iftop sysstat nvtop tmux curl wget git

CentOS/RHEL 系统：

sudo yum install -y epel-release
sudo yum install -y htop iotop iftop sysstat tmux curl wget git

安装 GPU 监控工具 nvtop：

sudo apt install -y nvtop

运行：

nvtop

三、Python 与 AI 框架环境优化

AI 编程常见框架包括 PyTorch、TensorFlow、JAX、Transformers、FastAPI、vLLM 等。环境不合理会造成大量性能损失。

1. 创建虚拟环境

推荐使用 Conda：

conda create -n ai-opt python=3.10 -y
conda activate ai-opt

或者使用 venv：

python3 -m venv ai-opt
source ai-opt/bin/activate

2. 安装 PyTorch GPU 版本

以 CUDA 12.1 为例：

pip install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu121

验证 CUDA 是否可用：

python -c "import torch; print(torch.__version__); print(torch.cuda.is_available()); print(torch.cuda.get_device_name(0))"

3. 安装常用 AI 开发库

pip install transformers datasets accelerate peft sentencepiece protobuf
pip install numpy pandas scikit-learn matplotlib tqdm
pip install fastapi uvicorn gunicorn
pip install onnx onnxruntime onnxruntime-gpu
pip install tensorrt

注意：TensorRT 的安装和 CUDA、cuDNN、驱动版本强相关，如果 pip 安装失败，应优先参考 NVIDIA 官方文档。

四、性能优化第一步：做 Profiling，而不是猜

很多性能问题不是模型慢，而是数据加载慢、日志打印慢、CPU 后处理慢、网络传输慢。如果不做 profiling，很容易优化错方向。

1. 使用 Python cProfile 分析

运行脚本：

python -m cProfile -o profile.out train.py

查看结果：

python -m pstats profile.out

进入交互界面后可以输入：

sort time
stats 20

含义是按耗时排序，显示前 20 个函数。

2. 使用 PyTorch Profiler

示例代码：

import torch
from torch.profiler import profile, record_function, ProfilerActivity

with profile(
    activities=[ProfilerActivity.CPU, ProfilerActivity.CUDA],
    record_shapes=True,
    profile_memory=True,
    with_stack=True
) as prof:
    with record_function("model_inference"):
        output = model(input_tensor)

print(prof.key_averages().table(sort_by="cuda_time_total", row_limit=20))

运行：

python profile_torch.py

导出 Chrome Trace：

prof.export_chrome_trace("trace.json")

然后在 Chrome 浏览器中打开：

chrome://tracing

加载 trace.json 即可查看执行时间线。

3. 使用 torch.utils.benchmark

import torch
import torch.utils.benchmark as benchmark

x = torch.randn(1024, 1024, device="cuda")
y = torch.randn(1024, 1024, device="cuda")

timer = benchmark.Timer(
    stmt="torch.matmul(x, y)",
    globals={"x": x, "y": y}
)

print(timer.timeit(100))

运行：

python benchmark_matmul.py

五、数据加载优化：很多训练慢都是 DataLoader 的锅

在 AI 训练中，GPU 利用率低并不一定是模型问题，常见原因是数据加载太慢。GPU 一直在等 CPU 准备数据，这会造成大量资源浪费。

1. PyTorch DataLoader 推荐配置

from torch.utils.data import DataLoader

train_loader = DataLoader(
    dataset,
    batch_size=64,
    shuffle=True,
    num_workers=8,
    pin_memory=True,
    persistent_workers=True,
    prefetch_factor=4
)

参数说明：

• num_workers：数据加载进程数量，通常设置为 CPU 核心数的 1/2 到 1 倍；
• pin_memory=True：加快 CPU 到 GPU 的数据拷贝；
• persistent_workers=True：避免每个 epoch 重复创建 worker；
• prefetch_factor：提前预取数据。

2. 找到合适的 num_workers

可以写一个简单脚本测试不同 worker 数量下的吞吐：

import time
from torch.utils.data import DataLoader

for workers in [0, 2, 4, 8, 16]:
    loader = DataLoader(
        dataset,
        batch_size=64,
        num_workers=workers,
        pin_memory=True,
        persistent_workers=True if workers > 0 else False
    )

    start = time.time()
    for i, batch in enumerate(loader):
        if i >= 100:
            break

    cost = time.time() - start
    print(f"num_workers={workers}, time={cost:.2f}s")

运行：

python test_dataloader_workers.py

3. 数据格式优化

如果数据是大量小文件，读取会非常慢。推荐将数据转换为更适合顺序读取的格式，例如：

• TFRecord；
• WebDataset；
• LMDB；
• Parquet；
• Arrow；
• Hugging Face datasets 缓存格式。

安装相关库：

pip install webdataset lmdb pyarrow datasets

将 CSV 转为 Parquet：

import pandas as pd

df = pd.read_csv("train.csv")
df.to_parquet("train.parquet", index=False)

运行：

python csv_to_parquet.py

读取 Parquet：

import pandas as pd

df = pd.read_parquet("train.parquet")

六、训练性能优化：让 GPU 真正跑满

1. 使用混合精度训练

混合精度可以显著减少显存占用，并提高训练速度。PyTorch 推荐使用 torch.cuda.amp。

import torch

scaler = torch.cuda.amp.GradScaler()

for batch in train_loader:
    inputs, labels = batch
    inputs = inputs.cuda(non_blocking=True)
    labels = labels.cuda(non_blocking=True)

    optimizer.zero_grad()

    with torch.cuda.amp.autocast():
        outputs = model(inputs)
        loss = criterion(outputs, labels)

    scaler.scale(loss).backward()
    scaler.step(optimizer)
    scaler.update()

运行训练：

python train.py

如果是 Hugging Face Transformers，可以直接开启 fp16：

python train.py \
  --model_name_or_path bert-base-chinese \
  --train_file train.json \
  --validation_file valid.json \
  --output_dir ./outputs \
  --per_device_train_batch_size 16 \
  --gradient_accumulation_steps 2 \
  --learning_rate 2e-5 \
  --num_train_epochs 3 \
  --fp16

如果 GPU 支持 BF16，例如 A100、H100、部分 RTX 40 系列，可以使用：

python train.py \
  --model_name_or_path bert-base-chinese \
  --output_dir ./outputs \
  --per_device_train_batch_size 16 \
  --bf16

2. 使用梯度累积扩大等效 Batch Size

当显存不足时，可以使用小 batch 加梯度累积。

等效 batch size 计算公式：

effective_batch_size = per_device_batch_size × gradient_accumulation_steps × GPU数量

命令示例：

python train.py \
  --per_device_train_batch_size 4 \
  --gradient_accumulation_steps 8 \
  --num_train_epochs 3 \
  --learning_rate 2e-5

3. 开启 cudnn benchmark

对于输入尺寸固定的视觉任务，可以开启：

import torch

torch.backends.cudnn.benchmark = True

如果需要完全可复现，则不要开启 benchmark，而是设置：

import torch
import random
import numpy as np

seed = 42
random.seed(seed)
np.random.seed(seed)
torch.manual_seed(seed)
torch.cuda.manual_seed_all(seed)

torch.backends.cudnn.deterministic = True
torch.backends.cudnn.benchmark = False

4. 使用 torch.compile 加速

PyTorch 2.x 提供了 torch.compile，可以对模型进行图编译优化。

import torch

model = model.cuda()
model = torch.compile(model)

运行：

python train_compile.py

也可以指定模式：

model = torch.compile(model, mode="max-autotune")

注意：首次运行会有编译开销，适合长时间训练或稳定输入形状的任务。

七、多 GPU 分布式训练优化

当单卡训练速度或显存不够时，可以使用多卡训练。最常见的是 PyTorch DDP。

1. 查看 GPU 数量

nvidia-smi -L

2. 使用 torchrun 启动多卡训练

单机 4 卡：

torchrun \
  --nproc_per_node=4 \
  train_ddp.py \
  --batch_size 32 \
  --epochs 10 \
  --lr 1e-4

指定使用某几张卡：

CUDA_VISIBLE_DEVICES=0,1 torchrun \
  --nproc_per_node=2 \
  train_ddp.py \
  --batch_size 32 \
  --epochs 10

3. NCCL 环境变量优化

多卡训练常用环境变量：

export NCCL_DEBUG=INFO
export NCCL_IB_DISABLE=0
export NCCL_P2P_DISABLE=0
export OMP_NUM_THREADS=8

启动命令：

NCCL_DEBUG=INFO OMP_NUM_THREADS=8 torchrun \
  --nproc_per_node=4 \
  train_ddp.py

如果遇到 NCCL 网络问题，可以临时禁用 IB：

NCCL_IB_DISABLE=1 torchrun \
  --nproc_per_node=4 \
  train_ddp.py

4. 使用 Accelerate 简化分布式训练

安装：

pip install accelerate

配置：

accelerate config

启动训练：

accelerate launch train.py

多卡启动：

accelerate launch \
  --num_processes 4 \
  train.py \
  --per_device_train_batch_size 8 \
  --gradient_accumulation_steps 4

八、显存优化：避免 OOM，提高可训练模型规模

1. 清理无用缓存

在合适位置释放缓存：

import torch
import gc

del loss
del outputs
gc.collect()
torch.cuda.empty_cache()

查看显存占用：

import torch

print(torch.cuda.memory_summary())

2. 使用梯度检查点

梯度检查点会用计算换显存，适合大模型训练。

model.gradient_checkpointing_enable()

Hugging Face 命令：

python train.py \
  --model_name_or_path big-model \
  --gradient_checkpointing \
  --per_device_train_batch_size 1 \
  --gradient_accumulation_steps 16

3. 使用 8-bit 优化器

安装 bitsandbytes：

pip install bitsandbytes

示例：

import bitsandbytes as bnb

optimizer = bnb.optim.AdamW8bit(model.parameters(), lr=2e-5)

4. 使用 LoRA/QLoRA 微调

安装：

pip install peft bitsandbytes accelerate transformers

示例命令：

python finetune_lora.py \
  --model_name_or_path Qwen/Qwen2.5-7B-Instruct \
  --train_file train.jsonl \
  --output_dir ./qwen-lora-output \
  --per_device_train_batch_size 1 \
  --gradient_accumulation_steps 16 \
  --learning_rate 2e-4 \
  --num_train_epochs 3 \
  --fp16 \
  --gradient_checkpointing

九、推理性能优化：从“能用”到“高并发可用”

1. 关闭梯度计算

推理时必须关闭梯度，否则会浪费显存和计算。

model.eval()

with torch.no_grad():
    outputs = model(inputs)

更推荐：

with torch.inference_mode():
    outputs = model(inputs)

2. 使用半精度推理

model = model.half().cuda()
inputs = inputs.half().cuda()

with torch.inference_mode():
    outputs = model(inputs)

对于 Transformers：

from transformers import AutoModelForCausalLM, AutoTokenizer
import torch

model = AutoModelForCausalLM.from_pretrained(
    "Qwen/Qwen2.5-7B-Instruct",
    torch_dtype=torch.float16,
    device_map="auto"
)

3. ONNX 导出与推理

安装：

pip install onnx onnxruntime onnxruntime-gpu

导出 ONNX：

import torch

dummy_input = torch.randn(1, 3, 224, 224).cuda()

torch.onnx.export(
    model,
    dummy_input,
    "model.onnx",
    input_names=["input"],
    output_names=["output"],
    dynamic_axes={
        "input": {0: "batch_size"},
        "output": {0: "batch_size"}
    },
    opset_version=17
)

运行：

python export_onnx.py

ONNX Runtime 推理：

import onnxruntime as ort
import numpy as np

session = ort.InferenceSession(
    "model.onnx",
    providers=["CUDAExecutionProvider", "CPUExecutionProvider"]
)

inputs = {
    session.get_inputs()[0].name: np.random.randn(1, 3, 224, 224).astype(np.float32)
}

outputs = session.run(None, inputs)
print(outputs[0].shape)

运行：

python infer_onnx.py

4. TensorRT 加速

如果模型结构适合 TensorRT，可以进一步加速。

ONNX 转 TensorRT：

trtexec \
  --onnx=model.onnx \
  --saveEngine=model_fp16.engine \
  --fp16 \
  --workspace=4096

动态 batch 示例：

trtexec \
  --onnx=model.onnx \
  --saveEngine=model_dynamic_fp16.engine \
  --fp16 \
  --minShapes=input:1x3x224x224 \
  --optShapes=input:8x3x224x224 \
  --maxShapes=input:32x3x224x224

测试 TensorRT 引擎性能：

trtexec \
  --loadEngine=model_fp16.engine \
  --warmUp=200 \
  --duration=60 \
  --iterations=1000

十、大语言模型推理优化：使用 vLLM

对于大语言模型，普通 Transformers 推理在高并发场景下效率较低。vLLM 通过 PagedAttention、连续批处理等机制，可以显著提升吞吐。

1. 安装 vLLM

pip install vllm

2. 启动 OpenAI 兼容服务

python -m vllm.entrypoints.openai.api_server \
  --model Qwen/Qwen2.5-7B-Instruct \
  --host 0.0.0.0 \
  --port 8000 \
  --dtype float16 \
  --gpu-memory-utilization 0.9 \
  --max-model-len 4096

3. 调用接口

curl http://localhost:8000/v1/chat/completions \
  -H "Content-Type: application/json" \
  -d '{
    "model": "Qwen/Qwen2.5-7B-Instruct",
    "messages": [
      {"role": "user", "content": "请用三句话解释什么是模型量化"}
    ],
    "temperature": 0.7,
    "max_tokens": 256
  }'

4. 限制最大并发和序列长度

如果显存容易爆，可以降低：

python -m vllm.entrypoints.openai.api_server \
  --model Qwen/Qwen2.5-7B-Instruct \
  --host 0.0.0.0 \
  --port 8000 \
  --dtype float16 \
  --gpu-memory-utilization 0.85 \
  --max-model-len 2048 \
  --max-num-seqs 64

十一、服务部署优化：FastAPI + Gunicorn

如果你的 AI 模型以 HTTP API 形式提供服务，需要优化 Web 服务层。

1. FastAPI 示例

from fastapi import FastAPI
import torch

app = FastAPI()

model = load_model()
model.eval()
model.cuda()

@app.post("/predict")
def predict(data: dict):
    inputs = preprocess(data)

    with torch.inference_mode():
        outputs = model(inputs)

    result = postprocess(outputs)
    return {"result": result}

启动开发服务：

uvicorn app:app --host 0.0.0.0 --port 8000

生产环境建议使用 Gunicorn：

gunicorn app:app \
  -k uvicorn.workers.UvicornWorker \
  --bind 0.0.0.0:8000 \
  --workers 1 \
  --timeout 120

注意：如果每个 worker 都会加载一份大模型，workers 数量不能盲目增加，否则显存会被复制占满。

2. 使用 Nginx 反向代理

安装：

sudo apt install -y nginx

配置文件：

sudo nano /etc/nginx/sites-available/ai-api

写入：

server {
    listen 80;
    server_name your-domain.com;

    client_max_body_size 50M;

    location / {
        proxy_pass http://127.0.0.1:8000;
        proxy_read_timeout 300;
        proxy_connect_timeout 300;
        proxy_send_timeout 300;
    }
}

启用配置：

sudo ln -s /etc/nginx/sites-available/ai-api /etc/nginx/sites-enabled/
sudo nginx -t
sudo systemctl restart nginx

十二、压测：优化必须用数据说话

性能优化一定要有压测数据，否则无法判断优化是否有效。

1. 使用 wrk 压测

安装：

sudo apt install -y wrk

GET 接口压测：

wrk -t4 -c100 -d60s http://127.0.0.1:8000/health

参数说明：

• -t4：4 个线程；
• -c100：100 个连接；
• -d60s：持续 60 秒。

2. POST 接口压测

创建 post.lua：

wrk.method = "POST"
wrk.body = '{"text":"请总结这段文本的主要内容"}'
wrk.headers["Content-Type"] = "application/json"

执行：

wrk -t4 -c50 -d60s -s post.lua http://127.0.0.1:8000/predict

3. 使用 hey 压测

安装：

go install github.com/rakyll/hey@latest

执行：

hey -n 1000 -c 50 http://127.0.0.1:8000/health

POST 压测：

hey \
  -n 1000 \
  -c 20 \
  -m POST \
  -H "Content-Type: application/json" \
  -d '{"text":"你好，请介绍一下人工智能"}' \
  http://127.0.0.1:8000/predict

十三、Docker 部署与性能优化

1. 编写 Dockerfile

FROM nvidia/cuda:12.1.1-cudnn8-runtime-ubuntu22.04

WORKDIR /app

RUN apt-get update && apt-get install -y \
    python3 python3-pip git curl \
    && rm -rf /var/lib/apt/lists/*

COPY requirements.txt .
RUN pip3 install --no-cache-dir -r requirements.txt

COPY . .

EXPOSE 8000

CMD ["uvicorn", "app:app", "--host", "0.0.0.0", "--port", "8000"]

2. 构建镜像

docker build -t ai-api:latest .

3. 使用 GPU 启动容器

docker run -d \
  --name ai-api \
  --gpus all \
  -p 8000:8000 \
  ai-api:latest

指定 GPU：

docker run -d \
  --name ai-api \
  --gpus '"device=0"' \
  -p 8000:8000 \
  ai-api:latest

4. 查看容器资源占用

docker stats ai-api

查看日志：

docker logs -f ai-api

进入容器：

docker exec -it ai-api bash

十四、常见性能问题与排查命令

问题 1：GPU 利用率很低

可能原因：

• DataLoader 太慢；
• batch size 太小；
• CPU 预处理太重；
• GPU 和 CPU 数据传输频繁；
• 模型本身太小，无法充分利用 GPU。

排查命令：

watch -n 1 nvidia-smi
htop
iotop

优化方向：

DataLoader(
    dataset,
    batch_size=128,
    num_workers=8,
    pin_memory=True,
    persistent_workers=True
)

问题 2：显存不足 OOM

排查：

nvidia-smi

查看 PyTorch 显存摘要：

print(torch.cuda.memory_summary())

优化方式：

• 减小 batch size；
• 开启 fp16/bf16；
• 开启梯度检查点；
• 使用 LoRA/QLoRA；
• 减小输入长度；
• 使用更小模型。

命令示例：

python train.py \
  --per_device_train_batch_size 1 \
  --gradient_accumulation_steps 16 \
  --fp16 \
  --gradient_checkpointing

问题 3：接口延迟很高

排查：

wrk -t4 -c20 -d30s http://127.0.0.1:8000/predict
top
nvidia-smi

优化方式：

• 模型常驻内存，不要每次请求重新加载；
• 使用 torch.inference_mode()；
• 批量推理；
• 使用 ONNX/TensorRT；
• LLM 使用 vLLM；
• 减少后处理逻辑；
• 使用缓存。

问题 4：CPU 占用过高

排查：

htop
pidstat -u 1

查看某个进程：

pidstat -p  1

优化方式：

• 减少 Python 循环；
• 使用 NumPy/PyTorch 向量化；
• 使用多进程；
• 将预处理前置；
• 使用缓存；
• C++/CUDA 扩展优化热点函数。

问题 5：磁盘 IO 成为瓶颈

排查：

iotop
iostat -x 1

优化方式：

• 使用 SSD/NVMe；
• 减少小文件；
• 数据转为 Parquet/LMDB/WebDataset；
• 增加数据缓存；
• 预加载常用数据。

十五、推荐的完整优化流程

一个成熟的 AI 项目性能优化流程可以按照以下步骤进行：

明确目标指标
    ↓
建立基准测试
    ↓
采集性能数据
    ↓
定位最大瓶颈
    ↓
只改一个变量
    ↓
重新压测和对比
    ↓
记录优化结果
    ↓
进入下一轮优化

建议每次优化都记录：

十六、一套可直接参考的训练启动命令

以下是一个较完整的单机多卡训练命令：

CUDA_VISIBLE_DEVICES=0,1,2,3 \
NCCL_DEBUG=INFO \
OMP_NUM_THREADS=8 \
torchrun \
  --nproc_per_node=4 \
  train.py \
  --model_name_or_path Qwen/Qwen2.5-7B-Instruct \
  --train_file ./data/train.jsonl \
  --validation_file ./data/valid.jsonl \
  --output_dir ./outputs/qwen-lora \
  --per_device_train_batch_size 1 \
  --gradient_accumulation_steps 16 \
  --learning_rate 2e-4 \
  --num_train_epochs 3 \
  --warmup_ratio 0.03 \
  --weight_decay 0.01 \
  --logging_steps 10 \
  --save_steps 500 \
  --eval_steps 500 \
  --fp16 \
  --gradient_checkpointing \
  --dataloader_num_workers 8

十七、一套可直接参考的推理服务启动命令

如果是大语言模型服务，推荐使用 vLLM：

CUDA_VISIBLE_DEVICES=0 \
python -m vllm.entrypoints.openai.api_server \
  --model Qwen/Qwen2.5-7B-Instruct \
  --host 0.0.0.0 \
  --port 8000 \
  --dtype float16 \
  --gpu-memory-utilization 0.9 \
  --max-model-len 4096 \
  --max-num-seqs 128

压测：

wrk -t4 -c50 -d60s http://127.0.0.1:8000/v1/models

接口调用：

curl http://127.0.0.1:8000/v1/chat/completions \
  -H "Content-Type: application/json" \
  -d '{
    "model": "Qwen/Qwen2.5-7B-Instruct",
    "messages": [
      {"role": "user", "content": "请解释 AI 编程中的性能优化方法"}
    ],
    "max_tokens": 512,
    "temperature": 0.7
  }'

十八、总结

AI 编程中的性能优化，不能只靠经验猜测，也不能只盯着模型本身。真正有效的优化，必须从完整链路出发：数据读取、CPU 预处理、GPU 计算、显存管理、模型结构、推理引擎、服务框架、网络传输、并发压测和系统监控，每一环都可能成为瓶颈。

如果是训练任务，优先关注 DataLoader、混合精度、batch size、梯度累积、分布式效率和显存优化。如果是推理任务，优先关注 模型常驻、关闭梯度、半精度、批处理、ONNX/TensorRT/vLLM、服务并发和压测数据。

最后记住一个原则：没有基准测试，就没有性能优化。 每次优化前先记录指标，每次修改只改变一个变量，然后用压测和 profiling 验证效果。这样才能避免“感觉变快了”的错觉，真正把 AI 项目优化到稳定、高效、低成本的生产级水平。