PyTorch报错：RuntimeError: Expected all tensors to be on the same device，如何解决？

Question

在PyTorch开发中，经常遇到`RuntimeError: Expected all tensors to be on the same device`错误。该问题通常发生在模型、输入张量或优化器涉及多个设备（如CPU和GPU）时，部分张量未被正确分配到同一设备。 **解决方法：** 1. 确保模型和输入数据在同一设备上。例如，将模型和张量移动到GPU（若有）： ```python device = torch.device("cuda" if torch.cuda.is_available() else "cpu") model.to(device) input_tensor = input_tensor.to(device) ``` 2. 检查模型参数、优化器状态是否一致。如果模型迁移到GPU，优化器中的参数也需要同步更新： ```python optimizer.state = {k: v.cuda() for k, v in optimizer.state.items()} ``` 3. 遍历复杂数据结构（如列表或字典），确保所有张量都在同一设备上。 通过以上步骤，可有效避免设备不匹配问题。

Answer 1

1. 问题概述

在PyTorch开发过程中，RuntimeError: Expected all tensors to be on the same device 是一个常见的错误。该问题通常发生在模型、输入张量或优化器涉及多个设备（如CPU和GPU）时，部分张量未被正确分配到同一设备。

以下是该问题的常见场景：

• 模型定义在GPU上，但输入数据仍位于CPU。
• 使用了多个GPU进行训练，但张量未正确同步。
• 优化器的状态字典未与模型参数保持一致。

2. 错误分析

当PyTorch执行计算图时，要求所有参与计算的张量必须位于同一设备上。如果模型的一部分在GPU上，而另一部分在CPU上，则会导致设备不匹配错误。

以下是一个简单的代码示例：

import torch
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model = MyModel().to(device)  # 模型加载到GPU
input_tensor = torch.randn(1, 3, 224, 224)  # 输入张量默认在CPU
output = model(input_tensor)  # 报错：设备不匹配

上述代码中，input_tensor 仍在CPU上，而模型已加载到GPU，因此触发错误。

3. 解决方案

为了解决设备不匹配问题，可以按照以下步骤逐步排查和修复：

3.1 确保模型和输入数据在同一设备上

将模型和张量显式地移动到相同的设备上：

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model.to(device)
input_tensor = input_tensor.to(device)

通过这种方式，可以确保模型和输入数据都在同一设备上。

3.2 同步优化器状态

如果模型从CPU迁移到GPU，优化器中的参数也需要同步更新：

for state in optimizer.state.values():
    for k, v in state.items():
        if isinstance(v, torch.Tensor):
            state[k] = v.to(device)

这段代码遍历优化器的状态字典，并将所有张量移动到目标设备。

3.3 遍历复杂数据结构

在实际项目中，输入数据可能是嵌套的列表或字典。此时需要递归地检查并移动所有张量：

def move_to_device(obj, device):
    if isinstance(obj, torch.Tensor):
        return obj.to(device)
    elif isinstance(obj, dict):
        return {k: move_to_device(v, device) for k, v in obj.items()}
    elif isinstance(obj, list):
        return [move_to_device(v, device) for v in obj]
    else:
        return obj

data = {"tensor1": torch.randn(3, 3), "tensor2": torch.randn(2, 2)}
data = move_to_device(data, device)

此函数支持递归处理嵌套数据结构，确保每个张量都正确移动到目标设备。

4. 流程图

以下是解决设备不匹配问题的流程图：

graph TD;
    A[开始] --> B[检查模型设备];
    B --> C{模型是否在GPU？};
    C --否--> D[将模型移动到GPU];
    C --是--> E[检查输入数据设备];
    E --> F{输入数据是否在GPU？};
    F --否--> G[将输入数据移动到GPU];
    F --是--> H[检查优化器状态];
    H --> I{优化器状态是否一致？};
    I --否--> J[同步优化器状态];
    I --是--> K[完成];

5. 总结性表格

以下是常见问题及解决方案的对比表：

问题类型	原因	解决方案
模型与输入设备不匹配	模型在GPU，输入在CPU	使用.to(device)统一设备
优化器状态不一致	模型迁移后优化器未更新	手动同步优化器状态
复杂数据结构问题	嵌套数据中部分张量未移动	递归遍历并移动张量