C++部署 pytorch 模型问题

[allegory]

目前采用的是使用 libtorch 加载 TorchScript 文件，然后 forward(input)完整这个过程，但问题来了：

1. torch script 文件是使用 python 代码将 pt 文件转换过来的，我已经在代码中指定了 cuda 设备编号为 4 ，部署上去后也在 c++代码中把所有的 tensor 、model 全部转移到 cuda:4 上去了，运行的时候还是会报错说有 hidden tensor 在 cuda:0 上，搞不懂
2. 我有 8 个 CUDA 设备，怎么样将它们全部利用起来跑我的数据？（又如何解决设备不一致的问题？）
3. 如何采用多线程跑 CUDA 任务

其实对于上面的问题不一定，得到答案不是最重要的，重要的是有没有前辈有这方面经验能多指教一下呢，torch 官方这方面的文档基本等于没有。（翻了又翻，文档全部是一些函数签名，没啥价值）

前辈们都是哪儿去找这方面的参考资料，可以分享一些吗，十分感谢

[allegory]

NVIDIA Driver Version: 545.23.08
CUDA Version: 12.3.107
PyTorch Version: 2.2.2+cu121
cuDNN Version: 8902
显卡驱动 环境啥的应该没问题

[ysc3839]

TorchScript 导出的不是只有 model 吗？怎么设置设备编号的？
正常应该是使用 torch::Device 来设置的吧，torch::jit::load 之后要 .to(device)

[rpWQTyfsAjMCKgPA]

GPU device 跟模型导出没有关系，跟 inference 时的设定有关。模型导出显然应该是设备无关的，你导出模型时写个 cuda:4 ，那模型拿到只有一个 GPU 的机器上难道就不能跑了？

[Muniesa]

第一个问题，不要在模型类代码中指定 cuda 设备，正确的做法是加载完模型后，使用 model.to("cuda:4")来指定。
2 和 3 建议用 triton inference server ，不需要你来处理多卡，当然你得先把代码按上面的改了。

[WoodsGao]

导出模型之前先 model.to("cpu")
load 的时候几张卡就 load 几回,每个 model 分别 to(每个 device)

[allegory]

感谢评论区各位大哥，不一一艾特了。
目前问题已经基本得到解决，设备不一致的问题主要是导出的时候用了 trace 而加载使用的时候用了 script ；至于后面两个问题，我也得到了解决，就是将数据多 to(device)几次放到了不同的 gpu 上。

[crazyliu]

看你原来那个贴子删掉了，只能在这里回复你了，需要内推可以看看我这个贴子 https://www.v2ex.com/t/1073632