pytorch - 如果我在 Cuda 上跟踪 PyTorch 网络，我可以在 CPU 上使用它吗？

Question

torch.jit.trace我在兼容 CUDA 的 GPU 服务器上跟踪了我的神经网络。当我在同一台服务器上重新加载该 Trace 时，我可以重新加载它并正常使用它。现在，当我将它下载到我的笔记本电脑上（用于快速测试）时，当我尝试加载跟踪时，我得到：

RuntimeError: Could not run 'aten::empty_strided' with arguments from the 'CUDA' backend. 'aten::empty_strided' is only available for these backends: [CPU, BackendSelect, Named, AutogradOther, AutogradCPU, AutogradCUDA, AutogradXLA, AutogradPrivateUse1, AutogradPrivateUse2, AutogradPrivateUse3, Tracer, Autocast, Batched, VmapMode].

我不能在跟踪上在 GPU 和 CPU 之间切换吗？还是有其他事情发生？

Answer 1

我有这个完全相同的问题。在我的模型中，我有一行代码导致了这种情况：

if torch.cuda.is_available():
    weight = weight.cuda()

如果您查看跟踪的官方文档（https://pytorch.org/docs/stable/generated/torch.jit.trace.html），您会看到

返回的 ScriptModule 将始终在任何输入上运行相同的跟踪图。当您的模块需要根据输入和/或模块状态运行不同的操作集时，这会产生一些重要的影响

因此，如果模型是在具有 GPU 的机器上跟踪的，则此操作将被记录下来，您甚至无法将模型加载到只有 CPU 的机器上。为了解决这个问题，删除了所有让你模型依赖于 CUDA 的东西。就我而言，就像删除上面的代码块一样简单。