2. 环境搭建：TensorFlow 2.x、TensorFlow Model Optimization Toolkit、TensorRT 8.x 安装与验证

好，咱们正式开始动手了。这一章说白了就是「把家伙事儿备齐」。我见过太多同学，代码写得挺溜，结果环境配了一星期，最后发现是版本不兼容——这种坑，我踩过不止一次。所以咱们稳一点，一步步来。

2.1 为什么这三样东西要一起装？

你想想看，咱们做模型优化，流程通常是这样的：

• 用 TensorFlow 2.x 训练模型
• 用 TensorFlow Model Optimization Toolkit 做剪枝和量化
• 最后用 TensorRT 8.x 做推理加速

这三者环环相扣。我个人习惯是先把环境锁死，再开始写代码。否则后面出了问题，你根本分不清是代码的锅还是环境的锅。

2.2 环境准备清单

嗯，先列个清单，咱们对着来：

组件	推荐版本	备注
Python	3.8 - 3.10	3.11 部分库还不支持
TensorFlow	2.12 或 2.13	2.14+ 有些 API 变了
TF Model Optimization	0.8.0	跟 TF 2.12 最搭
TensorRT	8.5 或 8.6	8.x 系列都行
CUDA	11.8	TensorRT 8.5 官方推荐
cuDNN	8.6	跟 CUDA 11.8 配对

2.3 安装 TensorFlow 2.x

这个最简单，一行命令搞定：

pip install tensorflow==2.12.0

如果你有 GPU，记得装 GPU 版本：

pip install tensorflow-gpu==2.12.0

验证一下：

python -c "import tensorflow as tf; print(tf.__version__)"

看到 2.12.0 就对了。我曾经遇到一个情况，装完发现是 CPU 版本，跑起来慢得离谱——后来才发现 pip 默认装的是 CPU 版，GPU 版得单独指定。

2.4 安装 TensorFlow Model Optimization Toolkit

这个库是咱们做剪枝和量化的核心工具。安装命令：

pip install tensorflow-model-optimization==0.8.0

验证：

python -c "import tensorflow_model_optimization as tfmot; print(tfmot.__version__)"

如果报错，八成是版本不匹配。我建议你直接用我上面给的版本号，别追新。

2.5 安装 TensorRT 8.x

这个稍微麻烦点。TensorRT 不能直接用 pip 装，得去 NVIDIA 官网下载。步骤是这样的：

1. 去 NVIDIA 官网下载 TensorRT 8.5 GA 版
2. 选择跟你的 CUDA 版本匹配的包（我用的 CUDA 11.8）
3. 解压后，把 lib 路径加到 LD_LIBRARY_PATH
4. 用 pip 安装 Python 绑定：pip install tensorrt-8.5.3.1-cp39-none-linux_x86_64.whl

验证：

python -c "import tensorrt as trt; print(trt.__version__)"

看到 8.5.3.1 就说明装好了。我记得第一次装 TensorRT 时，忘了加环境变量，结果 Python 里 import 成功但跑起来就崩——后来发现是动态库没找到。

2.6 联合验证

三个都装完后，咱们跑个完整的验证脚本：

import tensorflow as tf
import tensorflow_model_optimization as tfmot
import tensorrt as trt

print("TensorFlow version:", tf.__version__)
print("TF MOT version:", tfmot.__version__)
print("TensorRT version:", trt.__version__)

# 验证 GPU 是否可用
print("GPU available:", tf.config.list_physical_devices('GPU'))

# 验证 TensorRT 是否能加载
builder = trt.Builder(trt.Logger(trt.Logger.WARNING))
print("TensorRT builder created successfully")

如果所有版本都打印正常，GPU 也显示可用，那环境就齐活了。

2.7 常见问题与解决

嗯，这里列几个我实际遇到过的问题：

• 问题： TensorRT 报错 "Could not find libnvinfer.so"
  解决： 检查 LD_LIBRARY_PATH 是否包含 TensorRT 的 lib 目录
• 问题： TF Model Optimization 导入时报错 "cannot import name 'prune_low_magnitude'"
  解决： 版本不匹配，换成 0.8.0 试试
• 问题： TensorFlow 找不到 GPU
  解决： 检查 CUDA 和 cuDNN 版本，或者重装显卡驱动

说白了，环境搭建这件事，耐心比技术更重要。我当年刚开始做模型优化时，光配环境就花了一周。但一旦配好了，后面就顺了。

好，环境搭好了，下一章咱们就开始动手做剪枝了。