如何在C++中正确将PyTorch模型转换为半精度用于推理

你遇到的问题根源在于：手动重新注册参数时，新生成的半精度张量默认留在CPU，后续调用model.to(device)时，这些新注册的参数没有被模块的设备迁移逻辑正确同步到GPU，导致推理时输入与参数设备不匹配。

最简洁的解决方案 #

直接使用PyTorch C++中torch::jit::Module内置的to()方法，该方法与Python中model.half()/model.to(torch.float16)行为完全一致，会自动遍历所有参数、缓冲区完成 dtype 转换，同时支持后续正常的设备迁移：

#include <torch/script.h>

int main() {
    const std::string model_path = "your_model.pt";
    auto device = torch::kCUDA; // 或根据需求指定设备

    // 加载模型
    auto model = torch::jit::load(model_path);
    // 转换为半精度
    model = model.to(torch::kHalf);
    // 迁移到目标设备
    model = model.to(device);

    // 后续推理逻辑
    return 0;
}

也可以一步完成 dtype 转换与设备迁移：

model = model.to(device, torch::kHalf);

自定义转换逻辑的正确方式 #

如果需要自定义转换规则（比如仅转换特定层的参数），不要通过register_parameter重新注册参数，而是直接获取参数的可变引用修改其数据：

auto model = torch::jit::load(model_path);
model.apply([](torch::jit::Module& module) {
    // 处理当前模块的参数
    for (auto& param : module.named_parameters(false)) {
        if (param.value.is_floating_point()) {
            auto var = param.value.unwrap();
            var.data() = var.data().to(torch::kHalf);
        }
    }
    // 处理当前模块的缓冲区
    for (auto& buf : module.named_buffers(false)) {
        if (buf.value.is_floating_point()) {
            auto var = buf.value.unwrap();
            var.data() = var.data().to(torch::kHalf);
        }
    }
});
// 迁移到目标设备
model = model.to(device);

这种方式不会破坏模块内部的参数管理机制，后续调用to(device)时能正确同步所有参数到目标设备。

内容的提问来源于stack exchange，提问作者fei.sun