tensorflow,tf.map_fn在cpu上成功运行,但在gpu上失败

时间:2019-12-27 12:43:31

标签: python tensorflow

有两个代码块,一个可以成功运行,而另一个则失败。

代码是:

def success():
    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)
def failed():
    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("/gpu:0"):
        test = tf.map_fn(test_func, test_range, dtype=(tf.int32, tf.float32))

    sess.run(test)

失败的错误消息是

 [[node map/TensorArray_1 (defined at /github/KAM/seg_models/test.py:188) ]]Additional information about colocations:No node-device colocations were active during op 'map/TensorArray_1' creation.
No device assignments were active during op 'map/TensorArray_1' creation.

gpu和cuda的环境是正确的。 tf的版本是1.14.0 cuda的版本是10.0

0 个答案:

没有答案