我有一个深度学习(使用PyTorch(模型，其输出以字典格式给出。字典中有多个数组作为值，所有这些数组都在GPU内存中(torch.ttensors with device='cuda'(。有没有办法一次将字典中的每个数组移到主存中？
我目前的做法是写一个循环，将GPU数组重新写入变量并使用它们，但这似乎很无效。我的基本目标是以多线程的方式进一步处理输出，但由于这些阵列位于GPU内存上，多个线程无法同时访问它们。
如果有任何帮助或建议，我们将不胜感激！
您可以在输出字典out:上使用字典理解
out = {k: v.to(device='cpu', non_blocking=True) for k, v in out.items()}
如果out有一些元素不是张量，可以使用：
out = {k: v.to(device='cpu', non_blocking=True) if hasattr(v, 'to') else v for k, v in out.items()}