第2章环境准备：Python环境搭建、CUDA与cuDNN安装、Docker环境配置、PyTorch/TensorFlow框架安装

说实话，做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}")

📚 课程目录

第2章 环境准备：Python环境搭建、CUDA与cuDNN安装、Docker环境配置、PyTorch/TensorFlow框架安装

2.1 Python环境搭建

2.2 CUDA与cuDNN安装

2.3 Docker环境配置

2.4 PyTorch/TensorFlow框架安装

2.5 知识体系总览

第2章环境准备：Python环境搭建、CUDA与cuDNN安装、Docker环境配置、PyTorch/TensorFlow框架安装