第三章 环境搭建：PyTorch环境配置、CUDA与cuDNN安装、torch.quantization模块介绍、实验代码仓库初始化

说实话，做QAT量化感知训练，最怕的不是算法难，而是环境搭到一半崩了。我见过太多同学卡在CUDA版本不匹配上，一卡就是一整天。这一章，咱们就把这些坑提前填上。

3.1 PyTorch环境配置：选对版本，事半功倍

PyTorch的版本选择，我个人习惯是「追新不追旧，但别追最新」。什么意思？你想想看，量化这个模块在PyTorch 1.8之后才逐渐稳定下来。我建议直接用1.10以上版本，最好是1.13或2.0系列。

安装命令其实很简单，但有个细节要注意——别直接用pip install torch。为什么呢？因为这样装的是CPU版本。正确的姿势是去PyTorch官网，找到对应的安装命令。

# 我推荐的方式：用conda创建独立环境
conda create -n qat_env python=3.9
conda activate qat_env

# 安装PyTorch 1.13.1（CUDA 11.7版本）
pip install torch==1.13.1+cu117 torchvision==0.14.1+cu117 --extra-index-url https://download.pytorch.org/whl/cu117

# 验证安装
python -c "import torch; print(torch.__version__); print(torch.cuda.is_available())"

嗯，这里要注意：如果输出True，说明CUDA可用。如果输出False，别慌，大概率是CUDA驱动没装对。

3.2 CUDA与cuDNN安装：驱动、运行时、库，三层都要对

CUDA的安装，说白了就是三个东西：驱动（Driver）、运行时（Runtime）、深度神经网络库（cuDNN）。我刚开始做量化部署时，以为装个驱动就完事了，结果跑量化模型时各种报错。

为什么会这样？因为PyTorch的量化模块底层调用了cuDNN的算子。没有cuDNN，很多量化后的卷积操作就跑不起来。

安装步骤其实不复杂，我整理了一个清单：

1. 检查现有驱动：nvidia-smi查看驱动版本和最高支持的CUDA版本
2. 下载CUDA Toolkit：从NVIDIA官网下载，选择runfile或deb安装
3. 配置环境变量：在~/.bashrc中添加export PATH=/usr/local/cuda/bin:$PATH
4. 安装cuDNN：下载对应版本的cuDNN，解压后复制到CUDA目录
5. 验证安装：运行nvcc --version和cat /usr/local/cuda/include/cudnn_version.h | grep CUDNN_MAJOR

# cuDNN安装示例（以CUDA 11.7为例）
tar -xzvf cudnn-linux-x86_64-8.9.0.131_cuda11-archive.tar.xz
sudo cp cudnn-*-archive/include/cudnn*.h /usr/local/cuda/include 
sudo cp -P cudnn-*-archive/lib/libcudnn* /usr/local/cuda/lib64 
sudo chmod a+r /usr/local/cuda/include/cudnn*.h /usr/local/cuda/lib64/libcudnn*

3.3 torch.quantization模块介绍：三大核心组件

PyTorch的量化模块，说白了就三个东西：torch.quantization、torch.ao.quantization和torch.quantization.quantize_fx。嗯，名字有点绕，但功能很清晰。

我个人习惯用torch.ao.quantization，因为它是PyTorch 1.13之后主推的API。老版本的torch.quantization虽然也能用，但有些功能已经标记为废弃了。

核心组件就四个：

• QuantStub / DeQuantStub：量化入口和出口，负责把浮点张量转成量化张量，再转回来
• Observer：观察者，负责收集激活值和权重的统计信息（min/max/percentile等）
• FakeQuantize：伪量化模块，在训练中模拟量化误差
• QConfig：量化配置，告诉框架用什么Observer、量化位数等

# 一个最简单的量化配置示例
from torch.ao.quantization import QConfig, MinMaxObserver, PerChannelMinMaxObserver

# 我常用的配置：激活用MinMax，权重用PerChannel
qconfig = QConfig(
    activation=MinMaxObserver.with_args(dtype=torch.quint8, qscheme=torch.per_tensor_affine),
    weight=PerChannelMinMaxObserver.with_args(dtype=torch.qint8, qscheme=torch.per_channel_symmetric)
)

3.4 实验代码仓库初始化：从零搭一个可复现的实验框架

做量化实验，最怕的是跑完一个实验，过两天忘了参数怎么设的。我建议从一开始就搭一个规范的代码仓库。

我的习惯是这样的：

qat_from_scratch/
├── configs/          # 配置文件，YAML格式
│   ├── qat_config.yaml
│   └── ptq_config.yaml
├── data/             # 数据集存放
├── models/           # 模型定义
│   ├── __init__.py
│   ├── resnet.py
│   └── mobilenet.py
├── quant/            # 量化相关代码
│   ├── __init__.py
│   ├── observer.py   # 自定义Observer
│   ├── fake_quant.py # 自定义FakeQuantize
│   └── utils.py      # 量化工具函数
├── train.py          # 训练脚本
├── quantize.py       # 量化脚本
├── evaluate.py       # 评估脚本
└── requirements.txt  # 依赖清单

嗯，这里要注意：requirements.txt一定要写清楚版本号。我曾经因为没锁版本，三个月后复现实验时发现PyTorch升级了，量化结果对不上。

torch==1.13.1
torchvision==0.14.1
numpy==1.24.3
tqdm==4.65.0
tensorboard==2.13.0
pyyaml==6.0

初始化仓库时，我还会做一件事：写一个setup.sh脚本，把环境搭建、依赖安装、数据集下载全部自动化。这样换机器时，一行命令就能复现环境。

#!/bin/bash
# setup.sh - 一键环境搭建

# 创建conda环境
conda create -n qat_env python=3.9 -y
conda activate qat_env

# 安装依赖
pip install -r requirements.txt

# 下载示例数据集（以CIFAR-10为例）
python -c "
import torchvision
torchvision.datasets.CIFAR10(root='./data', train=True, download=True)
torchvision.datasets.CIFAR10(root='./data', train=False, download=True)
"

echo "环境搭建完成！"

3.5 本章知识体系总览

说了这么多，咱们用一张图把整个环境搭建的脉络理清楚。你想想看，从硬件驱动到Python库，再到实验框架，其实是一条完整的链路。

这张图把环境搭建的四个层次串起来了。从硬件驱动到CUDA库，再到PyTorch框架，最后到实验代码仓库。每一层都依赖下一层，所以排查问题时也要从底层往上查。

我记得有一次，量化模型在训练时突然报CUDA error: out of memory。我一开始以为是模型太大，折腾了半天才发现是cuDNN版本不对，导致某些算子用了更耗显存的实现。嗯，从那以后，我每次搭建环境都会先跑一遍torch.backends.cudnn.version()确认版本。

公众号:蓝海资料掘金营，微信deep3321