如何在CPU机器上测试PyTorch GPU代码并模拟CUDA设备校验张量迁移

当然可以实现这种模拟方案，核心思路是让所有张量的设备标识为CUDA，但实际仍在CPU上计算，一旦有未被转换的CPU张量参与运算就触发报错。下面提供两种实用的实现方式：

方法一：自定义伪CUDA张量类 #

通过继承PyTorch的Tensor类，创建一个设备标识为CUDA但底层用CPU计算的张量类型，同时拦截运算操作检查所有输入张量的类型：

import torch

class FakeCUDATensor(torch.Tensor):
    @classmethod
    def __torch_function__(cls, func, types, args=(), kwargs=None):
        # 遍历所有输入，检查是否存在未转换的普通CPU张量
        for arg in args:
            if isinstance(arg, torch.Tensor) and not isinstance(arg, FakeCUDATensor):
                raise RuntimeError(f"检测到未迁移至模拟CUDA设备的张量，当前设备: {arg.device}")
        
        # 执行原运算，结果包装为FakeCUDATensor
        result = super().__torch_function__(func, types, args, kwargs)
        if isinstance(result, torch.Tensor):
            return FakeCUDATensor(result.data, device=torch.device('cuda:0'))
        return result

# 转换工具函数：将普通CPU张量转为模拟CUDA张量
def to_fake_cuda(tensor):
    return FakeCUDATensor(tensor.data, device=torch.device('cuda:0'))

用法示例 #

# 正常场景：所有张量都转换为模拟CUDA类型
x = to_fake_cuda(torch.randn(2, 2))
y = to_fake_cuda(torch.randn(2, 2))
z = x + y  # 正常执行，无报错

# 错误场景：存在未转换的CPU张量
x = to_fake_cuda(torch.randn(2, 2))
y = torch.randn(2, 2)  # 未转换，设备为CPU
z = x + y  # 抛出RuntimeError，提示发现未迁移的张量

方法二：设备检查装饰器 #

通过装饰器包装模型的forward方法，在运算前后强制检查所有输入、输出张量的设备标识是否为指定的CUDA设备：

import torch
from functools import wraps

def enforce_fake_cuda(func):
    @wraps(func)
    def wrapper(*args, **kwargs):
        # 递归检查所有张量的设备
        def check_device(obj):
            if isinstance(obj, torch.Tensor):
                if obj.device != torch.device('cuda:0'):
                    raise RuntimeError(f"张量设备不匹配：预期cuda:0，实际为{obj.device}")
            elif isinstance(obj, (list, tuple)):
                for item in obj:
                    check_device(item)
            elif isinstance(obj, dict):
                for value in obj.values():
                    check_device(value)
        
        # 检查输入参数
        check_device(args)
        check_device(kwargs)
        
        # 执行原函数并检查输出
        result = func(*args, **kwargs)
        check_device(result)
        return result
    return wrapper

# 给模型的forward方法添加装饰器
class DemoModel(torch.nn.Module):
    def __init__(self):
        super().__init__()
        self.linear = torch.nn.Linear(3, 3)
    
    @enforce_fake_cuda
    def forward(self, x):
        return self.linear(x)

用法示例 #

model = DemoModel()
# 将模型参数的设备标识改为cuda:0（实际仍在CPU上存储）
for param in model.parameters():
    param.data = param.data.to(torch.device('cuda:0'))

# 正常场景：输入张量设备标识为cuda:0
x = torch.randn(4, 3).to(torch.device('cuda:0'))
output = model(x)  # 正常执行

# 错误场景：输入张量为CPU设备
x = torch.randn(4, 3)
output = model(x)  # 抛出RuntimeError，提示设备不匹配

注意事项 #

• 两种方案均无需真实GPU，完全在CPU上运行，适合本地开发调试
• 可覆盖模型前向传播、损失计算、张量运算等绝大多数场景
• 本质是通过设备标识校验替代真实CUDA运算，快速定位未指定设备的张量或模型参数

内容的提问来源于stack exchange，提问作者ikamen