现在,我尝试tf.nn.ctc_beam_search_decoder()
在 GPU 上使用。
但我有一个问题,它不使用 GPU。
我能够检查其他 tensorflow 函数(例如Reshape
等SigmoidGrad
)是否在 GPU 上运行。
但有些包括ctc_beam_search_decoder()
只在 CPU 上运行,而且ctc_beam_search_decoder()
速度很慢。
那么,我有两个问题。一、 TensorFlow2
不支持GPU?
其次,如果支持,你能告诉我如何实现或功能(或方法)吗?ctc_beam_search_decoder()
我在下面展示一个简单的例子。
程序代码。
import tensorflow as tf
from tensorflow.python.client import device_lib
tf.debugging.set_log_device_placement(True)
print(device_lib.list_local_devices())
inputs = tf.convert_to_tensor([
[0.1, 0.2, 0.3, 0.4, 0.5],
[0.2, 0.0, 0.3, 0.1, 0.1],
[0.2, 0.21, 0.3, 0.4, 0.1],
[0.2, 0.0, 0.6, 0.1, 0.5],
[0.2, 1.2, 0.3, 2.1, 0.1]])
inputs = tf.expand_dims(inputs, axis=1)
inputs_len = tf.convert_to_tensor([5])
decoded, _ = tf.nn.ctc_beam_search_decoder(inputs, inputs_len)
结果(标准输出)。
[name: "/device:CPU:0"
device_type: "CPU"
memory_limit: 268435456
locality {
}
incarnation: 714951449022474384
, name: "/device:XLA_CPU:0"
device_type: "XLA_CPU"
memory_limit: 17179869184
locality {
}
incarnation: 11733532016050292601
physical_device_desc: "device: XLA_CPU device"
, name: "/device:XLA_GPU:0"
device_type: "XLA_GPU"
memory_limit: 17179869184
locality {
}
incarnation: 394441871956590417
physical_device_desc: "device: XLA_GPU device"
, name: "/device:GPU:0"
device_type: "GPU"
memory_limit: 11150726272
locality {
bus_id: 1
links {
}
}
incarnation: 5917663253173554940
physical_device_desc: "device: 0, name: Tesla K80, pci bus id: 0000:00:04.0, compute capability: 3.7"
]
Executing op ExpandDims in device /job:localhost/replica:0/task:0/device:GPU:0
Executing op CTCBeamSearchDecoder in device /job:localhost/replica:0/task:0/device:CPU:0
Executing op StridedSlice in device /job:localhost/replica:0/task:0/device:GPU:0
忽略输入和输出数据并专注于正在使用的设备。
在这种情况下,ExpandDims
并StridedSlice
在 GPU 上执行。但CTCBeamSearchDecoder
并未在 GPU 上执行。