本文档旨在帮助解决在使用 TorchScript 模型时遇到的 "RuntimeError: Expected all tensors to be on the same device, but found at least two devices, cuda:0 and cpu" 错误。通过分析问题原因，提供详细的解决方案，并给出实际代码示例，确保模型在 CUDA 环境下正确运行，从而提高推理效率。
问题分析
这个错误通常发生在模型的一部分在 CUDA 设备上运行，而另一部分在 CPU 上运行的时候。PyTorch 要求所有参与运算的 Tensor 必须在同一个设备上。可能的原因包括：
- 模型内部某些层或操作显式地指定了 CPU 设备。
- 输入数据没有完全移动到 CUDA 设备上。
- TorchScript 在 tracing 过程中捕获了默认设备（CPU），导致模型在 CUDA 上运行时出现设备不一致。
解决方案
以下步骤可以帮助你解决这个问题：
1. 检查模型代码
首先，仔细检查你的模型代码，特别是模型定义和 forward 函数，查找任何显式指定 torch.device("cpu") 或类似操作的地方。如果发现，根据需要将其修改为使用 CUDA 设备。
2. 确保所有输入数据都在 CUDA 上
在将数据传递给模型之前，确保所有输入 Tensor 都已移动到 CUDA 设备。在 C++ 代码中，可以使用 .to(torch::kCUDA) 方法；在 Python 代码中，可以使用 .to(device) 方法，其中 device 是 CUDA 设备。
C++ 示例：
if (torch::cuda::is_available()) {
n_model = torch::jit::load("/home/lzh/Storage4/lzh/deepmodel/model_scripted.pt",torch::kCUDA);
std::cout << torch::cuda::device_count() << std::endl;
} else {
std::cerr << "No CUDA devices available, cannot move model to GPU." << std::endl;
}
torch::Tensor inputs = torch::from_blob(fre, {1, 4,300, 201}, torch::kFloat).to(torch::kCUDA);
std::cout << inputs.device() << std::endl;
textInput.input_ids.to(torch::kCUDA);
textInput.attention_mask.to(torch::kCUDA);
torch::Tensor out_tensor = n_model.forward({inputs,textInput.input_ids,textInput.attention_mask}).toTensor();Python 示例：
device=torch.device("cuda:0")
text = torch.ones((1,25))
text = text.long().to(device)
image = torch.ones((1,4,300,201)).to(device)
model = torch.jit.load('model_scripted.pt', map_location=torch.device('cuda'))
model.eval()
out = model(image,text,text)3. 在保存模型之前将模型移动到 CUDA 设备
这是最关键的一步。在 Python 中使用 torch.jit.trace 保存模型之前，先将整个模型移动到 CUDA 设备。这可以确保 TorchScript 在 tracing 过程中捕获 CUDA 设备信息。
Python 示例：
import torch
# 假设你的模型是 'model'
device = torch.device("cuda:0")
model.to(device) # 将模型移动到 CUDA 设备
# 创建一些示例输入，也移动到 CUDA 设备
image = torch.rand(1,4,300,201).to(device)
text1 = torch.rand(1,25).long().to(device)
text2 = torch.rand(1, 25).long().to(device)
# 使用 CUDA 设备上的输入 tracing 模型
traced_script_module = torch.jit.trace(model, (image,text1,text2))
# 保存 traced 模型
traced_script_module.save('model_scripted.pt')解释：
- model.to(device)：这行代码将模型的参数和缓冲区移动到指定的 CUDA 设备。
- 输入数据也需要移动到 CUDA 设备，确保 tracing 过程在 CUDA 上进行。
- torch.jit.trace(model, (image,text1,text2))：使用 CUDA 设备上的输入数据 tracing 模型。
4. 加载模型时指定 CUDA 设备 (C++)
在 C++ 中加载 TorchScript 模型时，确保指定 CUDA 设备。这可以通过在 torch::jit::load 函数中传递 torch::kCUDA 选项来实现。
C++ 示例：
torch::jit::Module n_model;
if (torch::cuda::is_available()) {
n_model = torch::jit::load("/path/to/model_scripted.pt", torch::kCUDA);
} else {
std::cerr << "No CUDA devices available, running on CPU." << std::endl;
n_model = torch::jit::load("/path/to/model_scripted.pt");
}总结
解决 "RuntimeError: Expected all tensors to be on the same device" 错误的关键在于确保模型和所有输入数据都在同一个设备上运行。在保存 TorchScript 模型之前将模型移动到 CUDA 设备，并在加载模型时指定 CUDA 设备，可以有效避免这个问题。
注意事项：
- 在调试过程中，可以使用 torch.cuda.current_device() 和 tensor.device 属性来检查当前设备和 Tensor 所在的设备。
- 如果你的模型包含多个子模块，请确保所有子模块的参数和缓冲区都已移动到 CUDA 设备。
通过遵循这些步骤，你应该能够成功解决 TorchScript 模型在 CUDA 环境中运行时的设备不一致问题，并充分利用 GPU 的加速能力。
