如何在GPU上使用多种返回类型定义tf.map_fn?

时间:2019-05-03 15:00:21

标签: tensorflow

import tensorflow as tf
import numpy as np

def test_func(i):
    return i, tf.add(tf.cast(i, tf.float32), 1.2)

test_range = tf.constant(np.arange(5))

sess = tf.Session()
sess.run(tf.global_variables_initializer())

with tf.device("/GPU:0"):
    test = tf.map_fn(test_func, test_range, dtype=(tf.int32, tf.float32))

sess.run(test)

运行上面的代码时,出现以下错误。我哪里错了?我的错误几乎与https://wedistribute.org/相似,但是由于有多种不同的返回类型,我遇到了一个奇怪的错误,如下所示:

InvalidArgumentError (see above for traceback): Cannot assign a device for operation map/TensorArray: Could not satisfy explicit device specification '' because the node node map/TensorArray (defined at test.py:13) having device No device assignments were active during op 'map/TensorArray' creation.  was colocated with a group of nodes that required incompatible device '/device:GPU:0'
Colocation Debug Info:
Colocation group had the following types and devices:
TensorArraySizeV3: CPU XLA_CPU XLA_GPU
TensorArrayV3: CPU XLA_CPU XLA_GPU
Range: GPU CPU XLA_CPU XLA_GPU
Const: GPU CPU XLA_CPU XLA_GPU
TensorArrayGatherV3: GPU CPU XLA_CPU XLA_GPU
Enter: CPU XLA_CPU XLA_GPU
TensorArrayReadV3: GPU CPU XLA_CPU XLA_GPU
TensorArrayScatterV3: CPU XLA_CPU XLA_GPU
TensorArrayWriteV3: CPU XLA_CPU XLA_GPU

Colocation members and user-requested devices:
  Const (Const)
  map/TensorArray (TensorArrayV3)
  map/TensorArrayUnstack/TensorArrayScatter/TensorArrayScatterV3 (TensorArrayScatterV3)
  map/TensorArray_1 (TensorArrayV3)
  map/while/TensorArrayReadV3/Enter (Enter) /device:GPU:0
  map/while/TensorArrayReadV3 (TensorArrayReadV3) /device:GPU:0
  map/while/TensorArrayWrite/TensorArrayWriteV3/Enter (Enter) /device:GPU:0
  map/while/TensorArrayWrite/TensorArrayWriteV3 (TensorArrayWriteV3) /device:GPU:0
  map/TensorArrayStack/TensorArraySizeV3 (TensorArraySizeV3)
  map/TensorArrayStack/range/start (Const)
  map/TensorArrayStack/range/delta (Const)
  map/TensorArrayStack/range (Range)
  map/TensorArrayStack/TensorArrayGatherV3 (TensorArrayGatherV3)

     [[node map/TensorArray (defined at test.py:13) ]]

No node-device colocations were active during op 'map/TensorArray' creation.
No device assignments were active during op 'map/TensorArray' creation.

1 个答案:

答案 0 :(得分:0)

我通过使用cpu解决了这个问题。我的电脑有GPU。

我不知道原因...

修改后的代码是

    import tensorflow as tf
    import numpy as np


    def test_func(i):
        return tf.cast(i, tf.int32), tf.add(tf.cast(i, tf.float32), 1.2)

    test_range = tf.constant(np.arange(5))

    sess = tf.Session()
    sess.run(tf.global_variables_initializer())

    with tf.device("/cpu:0"):
        test = tf.map_fn(test_func, test_range, dtype=(tf.int32, tf.float32))

    sess.run(test)