直接从TensorFlow访问PyTorch GPU矩阵

Question

我有一个用PyTorch编写的神经网络,它a在GPU上输出一些Tensor .我想继续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

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

a.storage()

一旦你有了存储空间，你就可以获得一个指向内存（CPU 或 GPU）的指针：

a.storage().data_ptr()

您可以检查它是否已固定

a.storage().is_pinned()

你可以固定它

a.storage().pin_memory()

我不熟悉 pyTorch 和 tensorflow 之间的接口，但我遇到了一个直接访问 GPU 中的 pytorch 张量的包（FAISS ）的示例。