PyTorch中torch.cuda.amp推理速度比普通FP32模式慢问题求助

• 硬件架构限制：你使用的Tesla P100属于帕斯卡架构，本身不支持Tensor Core，仅支持基础的FP16运算。autocast的核心加速逻辑依托Tensor Core完成混合精度的矩阵乘法、卷积运算，无该硬件能力时，开启autocast会额外增加FP32和FP16的动态类型转换开销，最终导致耗时翻倍。
• 版本适配问题：当前服务器的CUDA驱动支持的最高版本是11.2，但安装的PyTorch是基于CUDA10.1编译的，跨大版本的CUDA兼容性运行会引入额外的调度开销，进一步放大了autocast的性能损失。
• 测试负载不合理：测试用的batch size仅为1，整体计算量极小，类型转换的开销占比被无限放大，完全覆盖了FP16运算本身的速度优势。哪怕是支持Tensor Core的GPU，batch size过小的场景下也可能出现开启autocast变慢的情况。

• 更换支持Tensor Core的GPU：选择V100、T4、A10、A100等Volta及更新架构的GPU，这类硬件的Tensor Core可以将混合精度卷积、矩阵运算的速度提升数倍，完全抵消类型转换开销，体现autocast的加速效果。
• 对齐PyTorch与CUDA版本：卸载当前基于CUDA10.1编译的PyTorch，安装适配CUDA11.x版本的PyTorch，消除版本不兼容带来的额外开销。
• 增大测试batch size：将测试输入的batch size调整到32及以上，提升纯计算环节的占比，降低类型转换开销的权重，才能观测到混合精度的实际收益。
• 固定FP16推理（适合必须使用P100的场景）：如果只能使用P100显卡，可以放弃动态autocast，直接将模型和输入都转换为FP16格式执行推理，避免动态类型判断和转换的开销，参考代码：

model = torchvision.models.resnet18().half().to(device)
tensor = torch.rand(1,3,32,32).half().to(device)

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