PyTorch CUDA环境量化感知训练遇RuntimeError问题求助

针对AMP模式下的RuntimeError: expected scalar type Float but found Half问题 #

• 量化感知训练（QAT）的核心模块（如QuantStub、DeQuantStub及量化后的层）仅支持float32输入，但AMP会自动将张量转为float16，这是冲突根源。不能仅单独转换conv输出，要确保整个量化相关的前向路径强制在float32环境下运行：
  • 在QuantStub()输出后立刻添加x = x.float()，保证进入量化层的张量为float32；
  • 自定义前向函数时，将量化相关操作包裹在with torch.cuda.amp.autocast(enabled=False):上下文里，仅让非量化分支（如辅助损失计算）走AMP；
  • 检查优化器参数，确保所有参数都以float32注册，避免AMP自动转换优化器状态的dtype。

针对禁用AMP后的RuntimeError: Unsupported qscheme: per_channel_affine问题 #

• 先确认PyTorch版本：旧版本（1.10及以下）的CUDA量化不支持per_channel_affine方案，升级到PyTorch 1.11+可解决大部分兼容性问题；
• 手动修改量化配置，将层的qscheme替换为CUDA支持的类型：
  • 对Conv2d等层，在prepare_qat前设置：

    layer.qconfig = torch.ao.quantization.QConfig(
        activation=torch.ao.quantization.default_observer,
        weight=torch.ao.quantization.default_per_channel_weight_observer.with_args(qscheme=torch.per_tensor_affine)
    )
    

    把权重的qscheme从per_channel_affine换成per_tensor_affine；
  • 若使用默认qconfig（如default_qat_qconfig），替换为CUDA兼容的配置：qconfig = torch.ao.quantization.get_default_qat_qconfig('qnnpack')，qnnpack后端对CUDA量化方案的支持更完善。

额外注意事项 #

• 确保模型所有层都完成正确的量化准备：调用torch.ao.quantization.prepare_qat(model, inplace=True)前，需先给模型设置qconfig，且所有需要量化的层都被正确识别；
• 训练全程保证模型和数据都在CUDA设备上，避免中途切换到CPU，防止张量设备或 dtype 混乱。

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