直接从 TensorFlow 访问 PyTorch GPU 矩阵

Posted 2023-03-12

【中文标题】直接从 TensorFlow 访问 PyTorch GPU 矩阵

【英文标题】：Accessing PyTorch GPU matrix from TensorFlow directly

【发布时间】：2019-06-23 22:38:27

【问题描述】：

我有一个用 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()))

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

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

我尝试 2. 在上面使用 non_blocking=True 剪切的代码中，但我不确定它是否符合我的预期（即将其移动到固定内存）。

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

【参考方案1】：

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

a.storage()

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

a.storage().data_ptr()

你可以检查它是否被固定

a.storage().is_pinned()

你可以把它钉住

a.storage().pin_memory()

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

【讨论】：

在 pytorch 端，x.storage().data_ptr() 方法将数据指针作为 python 整数返回。不知道如何在 TF 端使用它。另外，要检查的一件事是计算是否在同一个 CUDA 流中执行，否则计算可能不同步。

@MatthijsDouze faiss 建议使用res.syncDefaultStreamCurrentDevice()，其中res 是当前计算资源的句柄。