在Windows 10上使用GPU提取图层输出时，TensorFlow 2.1.0有时会挂起冻结

如何解决在Windows 10上使用GPU提取图层输出时，TensorFlow 2.1.0有时会挂起冻结

我有一个自定义网络，经过训练后，我从网络中的图层中提取要素。我在具有1080Ti GPU的Windows 10上运行TF 2.1。

大多数时候，代码可以正常工作。有时，Keras会在调用下面的自定义函数或predict_on_batch后挂起。

#
# This function gets the layer output from a specific model
#
# verbose            Print debug info if set to True.
# model              TF model.
# x                  Input data (numpy array or [x1,x2,...,xN] for multi-input network).
# layer_number       The layer number,a non-negative integer.
# network_mode       Set to 0 to run the network in test mode (things like Dropout will be in effect).
#                    Set to 1 to run the network in train mode.
#                    Defaults to 0 (test mode).
# verbose            Print debug info if set to True.
#
def get_layer_output(model,x,layer_number,network_mode=0,verbose=False):

    if verbose:
        print("     >>> In get_layer_output")

    network_mode = tf.convert_to_tensor(network_mode)
    x = tf.convert_to_tensor(x)
    k_get_layer_output = K.function([model.layers[0].input,K.symbolic_learning_phase()],[model.layers[layer_number].output])
    if verbose:
        print("       >>> Calling k_get_layer_output")

    x_layer = k_get_layer_output([x,network_mode])[0]

    if verbose:
        print("     >>> Leaving get_layer_output")

    return x_layer

调用代码时，可以按预期方式工作99％的时间。但是，有时它会挂在这行上

layer = k_get_layer_output([x,network_mode])[0]

我从不离开该功能。

问题是间歇性的，但我确实将其追溯到TF / Keras被卡在上面的电话中。发生这种情况时，我的Python窗口将挂起并且代码冻结。

我花了大约一天的时间在网络上进行调查并找出问题所在。 另外，如果我用对predict_on_batch的调用替换了自定义函数，则会遇到同样的问题。

我怀疑是由GPU和TF引起的。 我将不胜感激。