4、环境搭建：PyTorch 量化工具包安装、ONNX Runtime 与 TensorRT 环境配置

说实话，环境搭建这事儿，看着简单，但坑是真不少。我见过太多同学在模型量化上花了大量精力，结果卡在环境配置上，一卡就是好几天。今天咱们就把这块彻底捋清楚。

4.1 PyTorch 量化工具包安装

PyTorch 从 1.8 版本开始，就把量化工具包直接集成到主框架里了。你不需要额外装什么包，但版本得选对。我个人习惯用 1.13 或 2.0 以上的版本，量化功能更稳定。

先确认你的 CUDA 版本：

nvidia-smi
# 或者
nvcc --version

然后安装对应版本的 PyTorch。我建议用 conda 创建独立环境，避免跟其他项目打架：

conda create -n qat_env python=3.9
conda activate qat_env

# CUDA 11.8 版本
pip install torch==2.0.1 torchvision==0.15.2 --index-url https://download.pytorch.org/whl/cu118

# 验证安装
python -c "import torch; print(torch.__version__); print(torch.backends.quantized.supported_engines)"

验证量化模块是否可用：

python -c "
import torch
from torch.quantization import quantize_dynamic, QuantStub, DeQuantStub
print('PyTorch 量化模块加载成功')
"

嗯，这里要注意。如果你看到类似 UserWarning: torch.quantization is deprecated 的警告，别慌。PyTorch 2.x 开始推荐用 torch.ao.quantization，但旧接口仍然能用。我个人建议新项目直接用 torch.ao.quantization。

4.2 ONNX Runtime 环境配置

ONNX Runtime 是微软开源的推理引擎，对量化模型支持得特别好。我一般在边缘设备上部署时，首选就是它。

安装很简单：

# CPU 版本
pip install onnxruntime

# GPU 版本（需要 CUDA）
pip install onnxruntime-gpu

# 安装 onnx 工具包
pip install onnx onnxconverter-common

验证安装：

python -c "
import onnxruntime
print('ONNX Runtime 版本:', onnxruntime.__version__)
print('可用执行提供者:', onnxruntime.get_available_providers())
"

如果你需要把 PyTorch 模型导出为 ONNX 格式，这里有个标准流程：

import torch
import torch.onnx

# 假设你有一个量化好的模型
model = torch.quantization.quantize_dynamic(
    original_model, {torch.nn.Linear}, dtype=torch.qint8
)

# 导出为 ONNX
dummy_input = torch.randn(1, 3, 224, 224)
torch.onnx.export(
    model,
    dummy_input,
    "model_quantized.onnx",
    opset_version=13,
    input_names=['input'],
    output_names=['output'],
    dynamic_axes={'input': {0: 'batch_size'}, 'output': {0: 'batch_size'}}
)

print("ONNX 模型导出成功")

为什么 opset_version 要设成 13？因为从 13 开始，ONNX 对量化操作的支持才比较完善。我试过用 11，结果导出后一堆算子不支持，白忙活一场。

4.3 TensorRT 环境配置

TensorRT 是 NVIDIA 的推理优化引擎，对量化模型的加速效果非常明显。说白了，如果你用的是 NVIDIA 的 GPU，TensorRT 几乎是必选项。

安装方式有两种：

1. 通过 pip 安装（推荐）：

pip install tensorrt

# 验证
python -c "import tensorrt as trt; print('TensorRT 版本:', trt.__version__)"

2. 通过 deb 包安装（更完整）：

去 NVIDIA 官网下载对应 CUDA 版本的 TensorRT 包。我建议用这种方式，因为会附带一些工具链，比如 trtexec 命令行工具，调试起来很方便。

安装完成后，测试一下能否正常构建引擎：

import tensorrt as trt
import numpy as np

# 创建 builder
logger = trt.Logger(trt.Logger.WARNING)
builder = trt.Builder(logger)

# 创建网络定义
network = builder.create_network(1 << int(trt.NetworkDefinitionCreationFlag.EXPLICIT_BATCH))

# 添加输入
input_tensor = network.add_input('input', trt.float32, (1, 3, 224, 224))

# 这里只是测试，实际使用时你会加载 ONNX 模型
print("TensorRT 环境配置成功，可以正常构建网络")

4.4 环境验证与常见问题

所有工具装完后，我建议跑一个完整的验证脚本：

import torch
import onnxruntime
import tensorrt as trt

print("=" * 50)
print("环境验证报告")
print("=" * 50)

# PyTorch
print(f"PyTorch 版本: {torch.__version__}")
print(f"CUDA 可用: {torch.cuda.is_available()}")
if torch.cuda.is_available():
    print(f"CUDA 版本: {torch.version.cuda}")
    print(f"GPU 型号: {torch.cuda.get_device_name(0)}")

# ONNX Runtime
print(f"ONNX Runtime 版本: {onnxruntime.__version__}")
print(f"可用提供者: {onnxruntime.get_available_providers()}")

# TensorRT
print(f"TensorRT 版本: {trt.__version__}")
print(f"TensorRT 构建版本: {trt.__version_build__}")

print("=" * 50)
print("所有环境配置完成！")
print("=" * 50)

最后说一句，环境搭建这事儿，一次配好，后面就省心了。我建议你把安装命令和版本号记下来，下次换机器或者重装系统时，直接复制粘贴就行。别问我为什么知道——我已经在不同机器上配过不下 20 次了。