从TensorFlow直接访问PyTorch GPU矩阵

Question

我已经写在PyTorch一个神经网络，输出GPU的一些张量a。我想继续处理a了高效TensorFlow层。

据我知道，要做到这一点的唯一方法是从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内存。

我试图在2上面的代码段使用non_blocking=True但我不知道它做什么，我希望（即其移动到固定的内存）。

理想情况下，我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之间的接口，但我碰到一个包（example）的FAISS来直接访问pytorch张量的GPU。