如何将TensorFlow SavedModel的float32权重转换为float16

• 加载SavedModel
  无需原模型代码，直接用TensorFlow加载已有的SavedModel：

  import tensorflow as tf
  
  model = tf.keras.models.load_model("./large_model")
  

• 递归转换所有层的权重
  考虑到模型可能存在嵌套层结构，通过递归遍历所有层，将每个权重张量转换为float16并更新层的权重：

  def convert_weights_to_float16(layer):
      # 转换当前层的可训练与非可训练权重
      if layer.weights:
          converted_weights = [tf.cast(weight, dtype=tf.float16) for weight in layer.weights]
          layer.set_weights([w.numpy() for w in converted_weights])
      
      # 递归处理子层（如嵌套的子模型、复合层）
      if hasattr(layer, "layers"):
          for sub_layer in layer.layers:
              convert_weights_to_float16(sub_layer)
  
  # 对整个模型执行权重转换
  convert_weights_to_float16(model)
  

• 保存转换后的模型
  将转换完成的模型保存到新目录，避免覆盖原始模型：

  model.save("./large_model_float16")
  

额外说明 #

• 转换后模型权重体积会减半（原3.6GB的权重文件将缩小至约1.8GB），达到减小体积的目的。
• 推理时TensorFlow会自动处理float16与输入数据的类型兼容问题，若需保证推理精度，可在推理阶段将输入数据转换为float16。
• 代码会自动跳过非浮点类型的权重（如整数类型的偏置参数），无需额外添加过滤逻辑。

内容的提问来源于stack exchange，提问作者Rasoul