2

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 之间切换吗?还是有其他事情发生?

4

1 回答 1

0

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

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

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

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

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

于 2022-01-19T14:43:02.167 回答