PyTorch中已训练torch.nn.Module模型的压缩与加速方法咨询

核心问题解答 #

PyTorch在评估模式下（通过model.eval()切换）不会自动压缩模型为等效函数，数据依然会按照原模型的层结构依次完成前向传播。不过此时像Dropout、BatchNorm这类训练时特殊的层会切换到评估行为（比如BatchNorm用训练时统计的均值方差，不再更新；Dropout关闭随机失活），但模型的整体结构不会被合并或简化。

提升评估速度与降低内存占用的方法 #

如果需要进一步优化推理效率，可采用以下几种实用方案：

• 模型量化：将模型的FP32精度权重转换为INT8甚至更低精度，能大幅减少内存占用并提升计算速度。PyTorch提供了动态量化（适合NLP模型）和静态量化（适合CV模型）两种方式，例如动态量化可以用torch.quantization.quantize_dynamic(model, {torch.nn.Linear}, dtype=torch.qint8)快速实现。
• 模型剪枝：移除模型中冗余的权重或通道，减少参数数量和计算量。可通过torch.nn.utils.prune模块实现，比如对卷积层进行结构化剪枝（剪掉整个通道），剪枝后还可对模型做少量微调恢复精度。
• 模型蒸馏：用训练好的大模型（教师模型）的输出作为监督，训练一个更小的模型（学生模型），得到性能接近但体积和计算量更小的模型，适合部署场景。
• TorchScript或ONNX转换：将模型转换为TorchScript（torch.jit.trace(model, example_input)）或ONNX格式，这类格式会自动进行层融合、算子优化等操作，同时还能配合TensorRT、ONNX Runtime等专用推理引擎进一步提升速度。
• 启用推理模式：使用torch.inference_mode()替代torch.no_grad()，它会彻底禁用梯度计算和反向传播相关的所有机制，比no_grad()更轻量，能减少内存开销。
• 自动混合精度推理：在GPU上推理时，启用torch.cuda.amp.autocast()，自动将部分算子转换为FP16精度，减少内存占用的同时几乎不损失精度，还能利用GPU的Tensor Core加速计算。

内容的提问来源于stack exchange，提问作者eilon toledano