Прямой доступ к матрице графического процессора PyTorch из TensorFlow - PullRequest
Новая
       33

Прямой доступ к матрице графического процессора PyTorch из TensorFlow

0 голосов
/ 30 января 2019

У меня есть нейронная сеть, написанная на PyTorch, которая выводит некоторый Tensor a на GPU.Я хотел бы продолжить обработку 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, но я предположил, что это невозможно?

1 Ответ

0 голосов
/ 06 февраля 2019

Я не знаком с тензорным потоком, но вы можете использовать pyTorch, чтобы раскрыть «внутренние» тензор.
Вы можете получить доступ к базовому хранилищу тензорного 

a.storage()

Как только у вас есть хранилище, вы можете получить указатель на память (либо CPU, либо GPU): 

a.storage().data_ptr()

Вы можете проверить, закреплено оно или нет 

a.storage().is_pinned()

И вы можете прикрепить его 

a.storage().pin_memory()

Я не знаком с интерфейсами между pyTorch и tenorflow, но я наткнулся на пример пакета ( FAISS ) напрямуюДоступ к тензорам Pytorch в GPU.

...