python - 如何将 tf.keras 与 bfloat16 一起使用
问题描述
我正在尝试使用混合精度让 tf.keras 模型在 TPU 上运行。我想知道如何使用 bfloat16 混合精度构建 keras 模型。是这样的吗?
with tf.contrib.tpu.bfloat16_scope():
inputs = tf.keras.layers.Input(shape=(2,), dtype=tf.bfloat16)
logits = tf.keras.layers.Dense(2)(inputs)
logits = tf.cast(logits, tf.float32)
model = tf.keras.models.Model(inputs=inputs, outputs=logits)
model.compile(optimizer=tf.keras.optimizers.Adam(.001),
loss='mean_absolute_error', metrics=[])
tpu_model = tf.contrib.tpu.keras_to_tpu_model(
model,
strategy=tf.contrib.tpu.TPUDistributionStrategy(
tf.contrib.cluster_resolver.TPUClusterResolver(tpu='my_tpu_name')
)
)
解决方案
您可以使用(计算和变量)使用下面显示的代码构建Keras
模型。bfloat16
Mixed Precision
float16
float32
tf.keras.mixed_precision.experimental.set_policy('infer_float32_vars')
model = tf.keras.Sequential([
tf.keras.layers.Inputs(input_shape=(2, ), dtype=tf.float16),
tf.keras.layers.Lambda(lambda x: tf.cast(x, 'float32')),
tf.keras.layers.Dense(10)])
model.compile(optimizer=tf.keras.optimizers.Adam(.001),
loss='mean_absolute_error', metrics=[])
model.fit(.............)
模型构建和训练后,我们可以使用以下步骤保存模型:
tf.keras.experimental.export_saved_model(model, path_to_save_model)
我们可以使用以下代码加载保存的混合精度 Keras 模型:
new_model = tf.keras.experimental.load_from_saved_model(path_to_save_model)
new_model.summary()
推荐阅读
- c - printf 或 scanf 问题(无法正确打印数字)
- http - 颤振:不记名令牌未发送到 API
- assembly - 我正在尝试将 C 代码转换为 MIPS 程序
- android - Google Pay 在 Xamarin Forms 上返回奇怪的错误
- javascript - Javascript 移除事件监听器
- javascript - 如何在 React Native 中从服务器获取数据之前渲染加载器
- sql - 分析几个表中的字段名
- angular - Angular-Slickgrid:并非所有过滤器都出现在标题行中
- c++ - vulkan-hpp 中奇怪的 C++ 语法
- javascript - 外部javascript可以在网站上访问什么。出于安全原因询问