在 tensorflow 模型训练中启用混合精度会降低而不是提高速度

Question

我正在使用 NVIDIA RTX-2060（带有图灵核心）进行深度学习模型训练。正如在线论坛上提到的，启用混合精度训练有助于图灵架构卡的训练速度比没有混合精度训练的速度更快。当我启用混合精度训练时，每步时间增加而不是减少。我无法理解为什么会发生这种情况，我非常感谢任何提出解决方案的人。我花了这么多钱买了这个gpu，如果我不能用它来更快地训练模型，那就没有用了。

代码：

import tensorflow as tf
def create_model():
    model = keras.Sequential([
        keras.layers.Flatten(input_shape=(32,32,3)),
        keras.layers.Dense(3000, activation='relu'),
        keras.layers.Dense(1000, activation='relu'),
        keras.layers.Dense(10, activation='sigmoid')
    ])

    
    model.compile(optimizer= 'SDG',
                 loss= 'categorical_crossentropy',
                 metrics=['accuracy'])
    
    return model

tf.keras.mixed_precision.set_global_policy('mixed_float16')

%%timeit -n1 -r1 # time required toexecute this cell once

model = create_model()
model.fit(X_train_scaled, y_train_categorical, epochs=50)

你必须知道的事情：

我已经成功安装了 cuda 和 cudnn 并且 tensorflow 可以检测到我的 gpu。

我已经安装了 tensorflow-gpu。

我正在使用 nvidia rtx 2060 gpu 在 ciphar10 数据集上训练我的模型。

我一直在使用 Jupyter Notebook 进行基准测试：Link

Answer 1

当您使用 CIFAR 数据集时，我认为您的最后一层激活应该softmax 而不是sigmoid，以及您的损失函数categorical_crossentropy。并确保它是float32。

keras.layers.Dense(10, activation='softmax', dtype=tf.float32)

---

您应该在开始时设置 混合精度 全局策略，就在 tf 导入之后。以下是在 GPU 上使用混合精度时的一些提示。来自Doc。

增加批量大小

如果不影响模型质量，请在使用 mixed precision 时尝试将 batch size 加倍运行。由于float16 张量使用了一半的内存，这通常可以让您将batch size 加倍而不会内存不足。增加 batch size 通常会增加训练吞吐量，即您的模型每秒可以运行的训练元素数。

确保使用 GPU 张量核心

现代 NVIDIA GPU 使用称为 张量核心 的特殊硬件单元，可以非常快速地将 float16 矩阵相乘。但是，Tensor Cores 要求张量的某些维度是 8 的倍数。

在下面的示例中，当且仅当它需要是 8 的倍数才能使用张量核心时，参数才是粗体。

• tf.keras.layers.Dense(units=64)
• tf.keras.layers.Conv2d(filters=48, kernel_size=7, stride=3)
• tf.keras.layers.LSTM(units=64)
• tf.keras.Model.fit(epochs=2, batch_size=128)

---

如果您正确遵循此程序，那么您应该获得使用mixed-precision 的优势。 Here 是 NVIDIA 的一篇好读物。

Answer 2

根据Tensorflow 的官方指南，要正确使用混合精度，模型末尾的sigmoid 激活应该是float32。因为我们设置了策略mixed_float16，所以激活的compute_dtype 是float16。因此，我们必须将这一层的策略覆盖为float32。

def create_model():
    model = keras.Sequential([
        keras.layers.Flatten(input_shape=(32, 32, 3)),
        keras.layers.Dense(3000, activation='relu'),
        keras.layers.Dense(1000, activation='relu'),
        # keras.layers.Dense(10, activation='sigmoid'), # NOTE: Replaced this line by two lines below
        keras.layers.Dense(10,),
        keras.layers.Activation('sigmoid', dtype='float32'),
    ])

    model.compile(optimizer='SGD',
                  loss='categorical_crossentropy',
                  metrics=['accuracy'])

    return model

把所有东西放在一起，我们就有了用于训练混合精度 CIFAR10 数据集的完整源代码：

import tensorflow as tf
from tensorflow import keras


def create_model():
    model = keras.Sequential([
        keras.layers.Flatten(input_shape=(32, 32, 3)),
        keras.layers.Dense(3000, activation='relu'),
        keras.layers.Dense(1000, activation='relu'),
        # keras.layers.Dense(10, activation='sigmoid'),
        keras.layers.Dense(10,),
        keras.layers.Activation('sigmoid', dtype='float32'),
    ])

    model.compile(optimizer='SGD',
                  loss='categorical_crossentropy',
                  metrics=['accuracy'])

    return model


tf.keras.mixed_precision.set_global_policy('mixed_float16')

(X_train, y_train), (X_test, y_test) = tf.keras.datasets.cifar10.load_data()

# There are 10 image classes
classes = ["airplane", "automobile", "bird", "cat", "deer", "dog", "frog", "horse", "ship", "truck"]

X_train_scaled = X_train / 255
X_test_scaled = X_test / 255

y_train_categorical = keras.utils.to_categorical(y_train, num_classes= 10, dtype='float')
y_test_categorical = keras.utils.to_categorical(y_test, num_classes= 10, dtype='float')

with tf.device('/GPU:0'):
    model = create_model()
    model.fit(X_train_scaled, y_train_categorical, epochs=50)

model.evaluate(X_test_scaled, y_test_categorical)

使用我的 GPU NVIDIA RTX2080，我使用混合精度比较了有（称为P1）和没有（称为P2）之间的性能，发现：

1. 训练时间：因为每一步的时间是四舍五入的，P1 和P2 的时间相同（~6ms），我比较了P1 的 50 个 epoch 的整体训练时间，即明显快于P2（464s 与 501s）。
2. 测试时间：P1 仍然比P2 快（3ms/step vs 4ms/step）
3. 测试性能（acc）：P1 优于 P2（~55.61% vs ~50.77%）