2. 环境准备
说实话,做Batch推理部署,环境准备这一步最容易被忽视。很多人上来就pip install,结果跑模型时各种报错,一查发现是CUDA版本不对。我刚开始带团队时也吃过这个亏,后来总结了一套标准流程,今天分享给你。
2.1 Python环境搭建
Python版本怎么选?我个人习惯用3.8或3.10。3.8兼容性最好,3.10性能稍好。别用3.12,有些老框架还没适配。
核心原则:每个项目一个虚拟环境,互不干扰。
我推荐用conda,原因很简单——它连CUDA toolkit都能一起管理。你想想看,要是用venv,还得手动装CUDA,多麻烦。
# 创建环境
conda create -n batch_infer python=3.10
# 激活环境
conda activate batch_infer
# 验证
python --version
pip --version
小技巧:把环境名写在项目根目录的.env文件里,团队协作时一目了然。
2.2 CUDA与cuDNN安装
CUDA版本和显卡驱动要匹配。怎么查?命令行跑一下:
nvidia-smi
看右上角的CUDA Version,那就是你驱动支持的最高版本。我建议装比这个低一个版本的CUDA,比如驱动支持12.2,你就装12.1,稳定。
cuDNN呢?说白了就是CUDA的加速库。装它有个坑——要注册NVIDIA账号才能下载。我曾经因为公司网络限制,折腾了半天才搞定。
注意:cuDNN版本必须和CUDA版本严格对应。8.x对应CUDA 11.x,9.x对应CUDA 12.x。别搞混了。
安装步骤:
- 下载CUDA runfile(别用deb,容易出依赖问题)
- 执行安装,选自定义安装,只勾选CUDA Toolkit
- 下载cuDNN,解压后复制到CUDA目录
# 验证CUDA
nvcc --version
# 验证cuDNN(在代码里测)
import torch
print(torch.cuda.is_available()) # 应该返回True
2.3 Docker环境配置
Docker是部署的神器。为什么?因为环境一致性。你在本地跑通了,到服务器上大概率也能跑通。
我一般用NVIDIA官方的基础镜像,省去装驱动的麻烦:
# 拉取PyTorch镜像
docker pull nvcr.io/nvidia/pytorch:23.12-py3
# 运行容器
docker run --gpus all -it --rm \
-v /path/to/your/code:/workspace \
nvcr.io/nvidia/pytorch:23.12-py3
关键参数说明:
--gpus all:让容器使用所有GPU-v:挂载本地代码目录--rm:退出后自动删除容器,避免堆积
嗯,这里要注意:Docker里跑CUDA程序,宿主机必须装好NVIDIA驱动和nvidia-docker2。我踩过这个坑——容器里nvcc --version正常,但跑模型就报错,最后发现是nvidia-container-runtime没装。
# 安装nvidia-docker2
distribution=$(. /etc/os-release;echo $ID$VERSION_ID)
curl -s -L https://nvidia.github.io/nvidia-docker/gpgkey | sudo apt-key add -
curl -s -L https://nvidia.github.io/nvidia-docker/$distribution/nvidia-docker.list | sudo tee /etc/apt/sources.list.d/nvidia-docker.list
sudo apt-get update
sudo apt-get install -y nvidia-docker2
sudo systemctl restart docker
2.4 PyTorch/TensorFlow框架安装
框架安装其实最简单,但有个原则:用官方推荐的pip命令,别自己瞎指定版本。
PyTorch安装:
# CUDA 12.1
pip install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu121
# CPU版本(调试用)
pip install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cpu
TensorFlow安装:
# GPU版本
pip install tensorflow[and-cuda]
# 验证
python -c "import tensorflow as tf; print(tf.config.list_physical_devices('GPU'))"
我的习惯:先装PyTorch,再装TensorFlow。因为PyTorch的CUDA依赖更干净,不容易冲突。
验证环境是否正常:
# PyTorch
python -c "import torch; print(torch.cuda.get_device_name(0))"
# TensorFlow
python -c "import tensorflow as tf; print(tf.test.is_gpu_available())"
如果输出正常,恭喜你,环境搭好了。如果报错,别慌——八成是CUDA版本不对,或者cuDNN没装对。
2.5 知识体系总览
下面这张图,是我整理的环境准备核心逻辑。你照着这个顺序来,基本不会出错。
这张图从下往上看:硬件→驱动→加速库→容器→框架。每一层都依赖下一层。我曾经跳过驱动检查直接装CUDA,结果浪费了半天时间。所以,老老实实按顺序来。
最后提醒:环境搭好后,跑一个简单的推理脚本测试。别等到部署时才发现问题。我习惯写个test_env.py,每次换机器先跑一遍。
# test_env.py
import torch
import tensorflow as tf
print(f"PyTorch CUDA: {torch.cuda.is_available()}")
print(f"TensorFlow GPU: {tf.test.is_gpu_available()}")
# 跑个小模型验证
x = torch.randn(100, 100).cuda()
y = torch.mm(x, x)
print(f"PyTorch 矩阵乘法测试通过: {y.shape}")
好了,环境准备就这些。说白了就是三板斧:装驱动、配CUDA、选框架。你按这个流程走,基本不会翻车。