在使用TorchScript模型时，可能会遇到 "RuntimeError: Expected all tensors to be on the same device, but found at least two devices, cuda:0 and cpu!" 错误。这个错误表明模型中的某些张量位于CPU上，而其他张量位于GPU上，导致操作无法顺利进行。根本原因是模型内部的某些操作可能在默认情况下创建了CPU张量，或者在加载模型后，某些张量没有正确地移动到GPU上。
解决方案：确保模型和输入数据在同一设备上
解决此问题的关键在于确保模型的所有参数以及所有输入数据都位于同一设备上，通常是CUDA设备（GPU）。以下步骤可以帮助你解决这个问题：
-
模型加载到GPU之前，先将模型移动到GPU:
在保存模型之前，先将模型移动到目标设备（例如CUDA），确保模型中的所有参数都位于GPU上。
device = torch.device("cuda:0") # 或者 "cpu" 如果你想在CPU上运行 model.to(device) -
在tracing之前，将输入移动到GPU:
在tracing模型时，使用的输入数据也应该位于与模型相同的设备上。
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) traced_script_module = torch.jit.trace(model, (image,text1,text2))
-
加载模型后，再次确认设备:
虽然在保存模型之前已经将模型移动到GPU，但在加载模型后，最好再次确认模型参数的设备。
model = torch.jit.load('model_scripted.pt', map_location=torch.device('cuda')) model.eval() # 设置为评估模式 for param in model.parameters(): if param.device.type == 'cuda': print('cuda') # 确认参数是否在cuda上 -
C++代码中的设备指定:
在C++代码中，确保加载模型时指定正确的设备，并且所有输入张量都已移动到该设备。
torch::Device device = torch::cuda::is_available() ? torch::kCUDA : torch::kCPU; torch::jit::Module n_model = torch::jit::load("/path/to/model_scripted.pt", device); torch::Tensor inputs = torch::from_blob(fre, {1, 4,300, 201}, torch::kFloat).to(device); textInput.input_ids = textInput.input_ids.to(device); textInput.attention_mask = textInput.attention_mask.to(device); torch::Tensor out_tensor = n_model.forward({inputs,textInput.input_ids,textInput.attention_mask}).toTensor();
注意事项：
- map_location 参数： 在使用 torch.jit.load 加载模型时，确保使用 map_location 参数将模型加载到正确的设备。 例如：torch.jit.load('model_scripted.pt', map_location=torch.device('cuda'))。
- 模型内部的设备指定： 检查模型代码，确保没有硬编码的设备指定。 如果模型内部有 torch.device('cpu') 这样的代码，可能会导致张量被创建在CPU上，从而引发设备不一致的错误。 尽量使用 device 变量来动态指定设备。
- 数据类型一致性： 确保输入数据的数据类型与模型期望的数据类型一致。 例如，如果模型期望 LongTensor 类型的输入，则确保输入数据是 LongTensor 类型。
- 评估模式： 在使用模型进行推理之前，务必将模型设置为评估模式：model.eval()。 这可以禁用 dropout 和 batch normalization 等训练时使用的层，从而提高推理效率和准确性。
总结：
解决 TorchScript 模型设备不一致问题的关键在于确保模型的所有组件（参数和输入数据）都位于同一设备上。通过在保存和加载模型时显式指定设备，并检查模型内部的设备指定，可以有效地避免此错误。在C++中使用模型时，也要确保将输入数据移动到与模型相同的设备。 遵循这些步骤可以确保模型在正确的设备上运行，并获得预期的结果。
