tensorflow混合精度训练

tensorflow 1.14后, 可以自动开启混合精度训练,在nvidia 某些设备上支持的很好,可以提速和节省显存,使用也很方便。

简介

看官方的介绍更好
Nvidia,auto_mix_precision

使用方法

在opt 后加上一句就好,
opt = tf.train.experimental.enable_mixed_precision_graph_rewrite(opt),

在终端会出现以下log,会自动将某些运算用fp16计算,并scale losslog

TensorFlow 2.5 中的混合精度训练可以通过使用 tf.keras.mixed_precision API 来实现。混合精度训练可以加快训练速度,同时减少 GPU 内存的使用。 要使用混合精度训练,需要先将模型的 dtype 设置为 float16,然后使用 tf.keras.mixed_precision.experimental.set_policy() 函数来设置混合精度策略。 以下是一个使用混合精度训练的示例: ```python import tensorflow as tf from tensorflow import keras from tensorflow.keras import layers from tensorflow.keras.mixed_precision import experimental as mixed_precision # 设置混合精度策略 policy = mixed_precision.Policy('mixed_float16') mixed_precision.set_policy(policy) # 构建模型 model = keras.Sequential([ layers.Dense(64, input_shape=(32,), activation='relu'), layers.Dense(10, activation='softmax') ], dtype='float16') # 编译模型 model.compile(optimizer='adam', loss='sparse_categorical_crossentropy', metrics=['accuracy']) # 加载数据集 (train_images, train_labels), (test_images, test_labels) = keras.datasets.mnist.load_data() train_images = train_images.reshape((60000, 28 * 28)).astype('float16') / 255 test_images = test_images.reshape((10000, 28 * 28)).astype('float16') / 255 # 训练模型 model.fit(train_images, train_labels, epochs=10, validation_data=(test_images, test_labels)) ``` 在上面的示例中,我们先设置混合精度策略,然后构建模型时将 dtype 设置为 float16,编译模型时指定优化器、损失函数和评估指标,最后使用 fit() 函数训练模型。
评论 4
添加红包

请填写红包祝福语或标题

红包个数最小为10个

红包金额最低5元

当前余额3.43前往充值 >
需支付:10.00
成就一亿技术人!
领取后你会自动成为博主和红包主的粉丝 规则
hope_wisdom
发出的红包
实付
使用余额支付
点击重新获取
扫码支付
钱包余额 0

抵扣说明:

1.余额是钱包充值的虚拟货币,按照1:1的比例进行支付金额的抵扣。
2.余额无法直接购买下载,可以购买VIP、付费专栏及课程。

余额充值