TorchScriptCUDA设备不一致解决方法

在使用TorchScript模型时，是否遇到过“RuntimeError: Expected all tensors to be on the same device, but found at least two devices, cuda:0 and cpu”的报错？本文针对这一常见问题，提供了一套详细的解决方案。首先，我们将深入分析错误原因，通常是由于模型部分在CUDA上运行，而另一部分在CPU上导致。其次，我们将提供一系列实用的解决步骤，包括检查模型代码、确保所有输入数据都在CUDA上、以及在保存模型前将模型移动到CUDA设备等。通过本文的指导，您将能够轻松解决TorchScript CUDA设备不一致的问题，确保模型在CUDA环境下高效运行，从而显著提升推理速度。文章附带C++和Python代码示例，方便您快速上手实践。

本文档旨在帮助解决在使用 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 必须在同一个设备上。可能的原因包括：

1. 模型内部某些层或操作显式地指定了 CPU 设备。
2. 输入数据没有完全移动到 CUDA 设备上。
3. 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 的加速能力。

好了，本文到此结束，带大家了解了《TorchScriptCUDA设备不一致解决方法》，希望本文对你有所帮助！关注golang学习网公众号，给大家分享更多文章知识！