tensorflow - 直接从 TensorFlow 访问 PyTorch GPU 矩阵

Question

我有一个用 PyTorch 编写的神经网络，它在 GPU 上输出一些张量 a。我想继续使用高效的 TensorFlow 层处理 a。

据我所知，执行此操作的唯一方法是将 a 从 GPU 内存移动到 CPU 内存，转换为 numpy，然后将其输入 TensorFlow。一个简化的例子:

import torch
import tensorflow as tf

# output of some neural network written in PyTorch
a = torch.ones((10, 10), dtype=torch.float32).cuda()

# move to CPU / pinned memory
c = a.to('cpu', non_blocking=True)

# setup TensorFlow stuff (only needs to happen once)
sess = tf.Session()
c_ph = tf.placeholder(tf.float32, shape=c.shape)
c_mean = tf.reduce_mean(c_ph)

# run TensorFlow
print(sess.run(c_mean, feed_dict={c_ph: c.numpy()}))

这可能有点牵强，但有没有办法做到这一点

1. a 永远不会离开 GPU 内存，或者
2. a 从 GPU 内存转到固定内存，再转到 GPU 内存。

我尝试在上面的代码中使用 non_blocking=True 进行 2. 操作，但我不确定它是否符合我的预期(即将其移动到固定内存)。

理想情况下，我的 TensorFlow 图将直接在 PyTorch 张量占用的内存上运行，但我认为这是不可能的？

Answer 1

我不熟悉 tensorflow ，但您可以使用 pyTorch 来公开张量的“内部结构”。
您可以访问底层storage张量的

a.storage()

一旦拥有存储空间，您就可以获得指向内存(CPU 或 GPU)的指针:

a.storage().data_ptr()

您可以检查它是否已固定

a.storage().is_pinned()

你可以固定它

a.storage().pin_memory()

我不熟悉pyTorch和tensorflow之间的接口(interface)，但我遇到了example直接访问 GPU 中的 pytorch 张量的包 ( FAISS )。