【30】使用混合精度(Mixed precision) 对训练产生的影响

Colab连结

一般我们在做机器学习任务时，在模型里计算的资料型态采用的是 float32 (即占用32的bits或4个bytes)，而 Nvidia 与 Baidu 一起发了一篇论文 Mixed Precision Training ，提出了一种训练方式叫混合精度(Mixed precision)，即训练过程中同时使用单精度(float32)和半精度(float16)。

在 Tensorflow 的文件中提到，若要开启双精度训练，只需要设定一个 global flag 即可。

mixed_precision.set_global_policy('mixed_float16')

如下范例，在开启此模式後，可以看到 dense 本身的权重值仍是 float32，但是计算和输出变成 float16。

inputs = keras.Input(shape=(784,), name='digits')
dense1 = layers.Dense(num_units, activation='relu', name='dense_1')
x = dense1(inputs)

print(dense1.dtype_policy)
print('x.dtype: %s' % x.dtype.name)
print('dense1.kernel.dtype: %s' % dense1.kernel.dtype.name)

产出:

<Policy "mixed_float16">
x.dtype: float16
dense1.kernel.dtype: float32

最後还有一个重要的一点，因为开启混合精度後，所有的计算预设会使用 float16，包括模型的最後一层 softmax，但我们会希望最後的 softmax 精度越高越好（越准确），因此我们要在模型最後的 softmax 中，设置 dtype='float32' 指定输出为float32。

(略)
x = layers.Dense(10, name='dense_logits')(x)
outputs = layers.Activation('softmax', dtype='float32', name='predictions')(x)

简单解释完上述概念後，我们今天要来实验的内容就是开启混合精度後，对我们的模型训练有什麽影响？

实验一：用混合精度训练 mnist

mixed_precision.set_global_policy(mix_policy)

model = tf.keras.Sequential()
model.add(tf.keras.layers.Conv2D(32, [3, 3], activation='relu', input_shape=(28,28,1)))
model.add(tf.keras.layers.Conv2D(64, [3, 3], activation='relu'))
model.add(tf.keras.layers.MaxPooling2D(pool_size=(2, 2)))
model.add(tf.keras.layers.Dropout(0.25))
model.add(tf.keras.layers.Flatten())
model.add(tf.keras.layers.Dense(128, activation='relu'))
model.add(tf.keras.layers.Dropout(0.5))
model.add(tf.keras.layers.Dense(10))
model.add(tf.keras.layers.Softmax(dtype='float32'))  # float32

print(f'{model.layers[1].name}:{model.layers[1].dtype_policy}')

model.compile(
    optimizer=tf.keras.optimizers.SGD(LR),
    loss=tf.keras.losses.SparseCategoricalCrossentropy(from_logits=False),
    metrics=[tf.keras.metrics.SparseCategoricalAccuracy()],
)

history = model.fit(
    ds_train_tf,
    epochs=EPOCHS,
    validation_data=ds_test_tf,
)

产出：

loss: 0.0475 - sparse_categorical_accuracy: 0.9851 - val_loss: 0.0311 - val_sparse_categorical_accuracy: 0.9899

可以看到每次 epoch 平均花费8秒多。

实验二：用单精度训练 mnist

mixed_precision.set_global_policy(f32_policy)

model = tf.keras.Sequential()
model.add(tf.keras.layers.Conv2D(32, [3, 3], activation='relu', input_shape=(28,28,1)))
model.add(tf.keras.layers.Conv2D(64, [3, 3], activation='relu'))
model.add(tf.keras.layers.MaxPooling2D(pool_size=(2, 2)))
model.add(tf.keras.layers.Dropout(0.25))
model.add(tf.keras.layers.Flatten())
model.add(tf.keras.layers.Dense(128, activation='relu'))
model.add(tf.keras.layers.Dropout(0.5))
model.add(tf.keras.layers.Dense(10))
model.add(tf.keras.layers.Softmax(dtype='float32'))  # float32

print(f'{model.layers[1].name}:{model.layers[1].dtype_policy}')

model.compile(
    optimizer=tf.keras.optimizers.SGD(LR),
    loss=tf.keras.losses.SparseCategoricalCrossentropy(from_logits=False),
    metrics=[tf.keras.metrics.SparseCategoricalAccuracy()],
)

history = model.fit(
    ds_train_tf,
    epochs=EPOCHS,
    validation_data=ds_test_tf,
)

产出：

loss: 0.0496 - sparse_categorical_accuracy: 0.9845 - val_loss: 0.0310 - val_sparse_categorical_accuracy: 0.9893

每次 epoch 花费7秒左右，跟预期的不太一样，使用混合精度结果训练比较慢。

实验三：用混合精度训练 oxford_flowers102 (训练程序码都相同，以下皆省略)

loss: 4.2232e-04 - sparse_categorical_accuracy: 1.0000 - val_loss: 0.4107 - val_sparse_categorical_accuracy: 0.8853

每次 epoch 花费10秒左右

实验四：用单精度训练 oxford_flowers102

loss: 4.8664e-04 - sparse_categorical_accuracy: 1.0000 - val_loss: 0.4102 - val_sparse_categorical_accuracy: 0.8843

每次 epoch 花费13秒，混合精度胜！

实验五：用混合精度训练 cifar10

loss: 0.0448 - sparse_categorical_accuracy: 0.9843 - val_loss: 0.4190 - val_sparse_categorical_accuracy: 0.9001

每次 epoch 花费在417秒左右

实验六：用单精度训练 cifar10

loss: 0.0419 - sparse_categorical_accuracy: 0.9859 - val_loss: 0.3610 - val_sparse_categorical_accuracy: 0.9091

每次 epoch 花费秒数在566~574区间，仍是混合精度胜！

综观以上结果，我们发现两者准确度都差不多，但是在模型比较复杂的状况下，使用混合精度节省了不少时间，因此在使用 GPU 的状况下，都非常推荐将混合精度开启。