2

我正在使用 keras 混合精度 API 以适应 GPU 中的网络。通常在我的代码中，这看起来像这样。MWE 将是：

from tensorflow.keras.mixed_precision import experimental as mixed_precision

use_mixed_precision = True

if use_mixed_precision:
  policy_type = 'mixed_float16'
else:
  policy_type = 'float32'
policy = mixed_precision.Policy(policy_type)
mixed_precision.set_policy(policy)

这似乎达到了预期的效果，因为当我训练我的模型并使用 TensorBoard 回调对其进行分析时，我的大部分操作都以半精度运行，其中一些正在使用 TensorCore（我有一个具有计算能力的 GPU 7.0 以上）。

然而，Conv2DBackpropFilter没有使用 TensorCore，即使根据 TensorBoard 信息它有资格使用它。

TensorCore 操作

我还没有整个事情的最小可重现示例，如果需要，我可以处理它，但我想首先知道这是否是预期的行为，或者是否有一些已知的陷阱，因为我找不到任何网上资料。

编辑

我有一个 MRE，它的行为不同，但问题相同：为什么不使用 TensorCore（所有需要的维度都是 8 的倍数）？

import tensorflow as tf
from tensorflow.keras.mixed_precision import experimental as mixed_precision

use_mixed_precision = True

if use_mixed_precision:
    policy_type = 'mixed_float16'
else:
    policy_type = 'float32'
policy = mixed_precision.Policy(policy_type)
mixed_precision.set_policy(policy)

nf = 8
model = tf.keras.models.Sequential([
    tf.keras.layers.Conv2D(filters=nf, kernel_size=3, padding='same'),
    tf.keras.layers.Conv2D(filters=nf, kernel_size=3, padding='same'),
    tf.keras.layers.Conv2D(filters=nf, kernel_size=3, padding='same'),
])
model.compile(loss='mse', optimizer='sgd')

bs = 8
inputs = tf.random.normal([bs, 32, 32, 1])
outputs = tf.random.normal([bs, 32, 32, nf])

tboard_cback = tf.keras.callbacks.TensorBoard(
    profile_batch='5, 10',
    log_dir='logs',
    histogram_freq=0,
    write_graph=False,
    write_images=False,
)

model.fit(inputs, outputs, callbacks=[tboard_cback], epochs=15)

在这个 MRE 中，64.2% 的操作时间花费在半精度上，这意味着半精度确实正在发生。在我的日志中，我还检查了计算能力：

NFO:tensorflow:Mixed precision compatibility check (mixed_float16): OK
Your GPU will likely run quickly with dtype policy mixed_float16 as it has compute capability of at least 7.0. Your GPU: Tesla V100-SXM2-32GB, compute capability 7.0

然而，没有任何操作（这次不仅仅是Conv2DBackpropFilter）使用 TensorCore 运行。我不明白为什么。

tensorcore_for_mre

4

0 回答 0