解决 TorchScript 模型 CUDA 设备不一致问题：教程与实践-Python教程-PHP中文网

解决 TorchScript 模型 CUDA 设备不一致问题：教程与实践

霞舞

发布： 2025-07-31 19:02:11

原创

990人浏览过

解决 torchscript 模型 cuda 设备不一致问题：教程与实践

本文档旨在帮助解决在使用 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 设备信息。

文心大模型

百度飞桨-文心大模型 ERNIE 3.0 文本理解与创作

查看详情

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 设备，可以有效避免这个问题。

注意事项：