Python 3.x 我可以使用双变量方法异步预取tensorflow（1.4）中的GPU吗？

据我所知，tensorflow 1.4中CPU/GPU之间仍然没有异步预取数据

我试图自己编写这个功能，作为理解的练习

以下代码尝试实现此过程：

下一批存储在

data\u var

变量中

处理一些从变量

data\u var

获取批次的操作（也称为训练模型），在本例中，只需执行

myop

，这取决于

data\u var

prefetch

OP将下一批数据从数据集对象（基于CPU）分配到

temp\u var

基于gpu的变量中

使用

tf.control\u dependencies（…）

等待

myop

和

prefetch\u op

完成，然后将

temp\u var

分配给

data\u var

这似乎不起作用。TF探查器显示，

myop

并不像所希望的那样与

MEMCPYHtoD

进程进行异步处理

我原以为这两个操作，

myop

和

prefetch\u op

将异步运行，因为它们之间没有依赖关系

下面是我用来运行这个测试的代码。它将独立运行

import tensorflow as tf
from tensorflow.python.client import timeline
import numpy as np
import os

sz = 2000
x = np.random.rand(sz, sz)

def gen():
  yield x

# Dataset
ds = tf.data.Dataset.from_generator(generator=gen, output_types=tf.float64)
ds = ds.repeat()
ds = ds.prefetch(2)
iterator = ds.make_one_shot_iterator()
next_element = iterator.get_next()

# Prefetch to GPU OPs - this is exepected to happen asynchronously
temp_var = tf.Variable(np.zeros((sz, sz)), name='temp_var', dtype=tf.float64, expected_shape=(sz, sz), trainable=False)
data_var = tf.Variable(np.zeros((sz, sz)), name='data_var', dtype=tf.float64, expected_shape=(sz, sz), trainable=False)
prefetch_op = tf.assign(temp_var, next_element)

# Trivial math operation for timing purposes
myop = tf.sqrt(data_var, name='myop')

# Final prefetch to GPU operation, copy data from temp_var to data_var
with tf.control_dependencies((myop, prefetch_op)):
  assign_op = tf.assign(data_var, temp_var)

# Open session, initialize, and run 1 iteration to warm the prefetch buffer
sess = tf.Session()
sess.run(tf.global_variables_initializer())
sess.run((myop, assign_op))

# Main sess.run with profiling on
tf_options_profiler_on = tf.RunOptions(trace_level=tf.RunOptions.FULL_TRACE)
tf_run_metadata = tf.RunMetadata()
graph_result = sess.run((myop, assign_op), options=tf_options_profiler_on, run_metadata=tf_run_metadata)

# Write profile data
chrome_trace = timeline.Timeline(tf_run_metadata.step_stats).generate_chrome_trace_format()
os.makedirs('/tmp/profile', exist_ok=True)
with open('/tmp/profile/tf_profile_step.json', 'w') as f:
  f.write(chrome_trace)

print(graph_result)
print('Writing profiler output to /tmp/profile')

---

在tensorflow 1.7中，数据集API现在具有

预取\u到\u设备

文件：

Github讨论：

在上面的Github讨论（现已结束）中，似乎进一步提到了另一个选项，称为

多设备迭代器

。