MacOS下PyInstaller打包TensorFlow程序，保存模型遇TensorShapeProto错误

问题背景 #

编写的TensorFlow模型保存/加载测试脚本在系统Python环境运行正常，但通过PyInstaller打包后，MacOS环境下执行时保存模型环节报错，Windows环境无此问题。若提前在外部保存模型，仅通过打包程序加载时会报AttributeError: as_proto，推测二者根因一致。

测试脚本 #

import tempfile
from dataclasses import dataclass
import cv2
import tensorflow as tf


@dataclass
class ColorFilterModel(tf.Module):
    rgb_filter_strengths = (0., 1., 1.)
    def compute_mahalonabis_sq_heatmap(self, image):
        return tf.reduce_sum(tf.cast(image, tf.float32) * tf.constant(self.rgb_filter_strengths), axis=-1)

def test_basic_save_and_load_model():
    print("Testing basic model save/load")
    model = ColorFilterModel()
    numpy_image = cv2.imread(tf.keras.utils.get_file('basalt_canyon.jpg', "https://raw.githubusercontent.com/petered/data/master/images/basalt_canyon.jpg"))
    print(f"Image shape: {numpy_image.shape}")
    with tempfile.TemporaryDirectory() as tmpdirname:
        concrete_func = tf.function(model.compute_mahalonabis_sq_heatmap, input_signature=[tf.TensorSpec(shape=(None, None, 3), dtype=tf.uint8)]).get_concrete_function()
        result_original = concrete_func(numpy_image)
        print(f"Result shape: {result_original.shape}")
        tf.saved_model.save(model, tmpdirname, signatures={'mahal': concrete_func})  # 报错行
        loaded_func = tf.saved_model.load(tmpdirname).signatures['mahal']
        result_loaded = loaded_func(image=numpy_image)['output_0']
        assert tf.reduce_all(tf.equal(result_original, result_loaded))
    print("Passed test_basic_save_and_load_model")


if __name__ == "__main__":
    test_basic_save_and_load_model()

打包后报错信息 #

首先出现警告：

WARNING:tensorflow:AutoGraph is not available in this environment: functions lack code information. This is typical of some environments like the interactive Python shell. See https://github.com/tensorflow/tensorflow/blob/master/tensorflow/python/autograph/g3doc/reference/limitations.md#access-to-source-code for more information.

随后抛出错误：

Testing basic model save/load
Image shape: (1500, 2000, 3)
2023-05-24 17:20:13.912138: W tensorflow/core/platform/profile_utils/cpu_utils.cc:128] Failed to get CPU frequency: 0 Hz
Result shape: (1500, 2000)
Traceback (most recent call last):
  File "main.py", line 26, in <module>
    test_basic_save_and_load_model()
  File "video_scanner/detection_utils/test_standalone_model_io.py", line 22, in test_basic_save_and_load_model
    tf.saved_model.save(model, tmpdirname, signatures={'mahal': concrete_func})
  File "tensorflow/python/saved_model/save.py", line 1232, in save
  File "tensorflow/python/saved_model/save.py", line 1268, in save_and_return_nodes
  File "tensorflow/python/saved_model/save.py", line 1441, in _build_meta_graph
  File "tensorflow/python/saved_model/save.py", line 1396, in _build_meta_graph_impl
  File "tensorflow/python/saved_model/save.py", line 794, in _fill_meta_graph_def
  File "tensorflow/python/saved_model/save.py", line 607, in _generate_signatures
  File "tensorflow/python/saved_model/save.py", line 474, in _tensor_dict_to_tensorinfo
  File "tensorflow/python/saved_model/save.py", line 475, in <dictcomp>
  File "tensorflow/python/saved_model/utils_impl.py", line 78, in build_tensor_info_internal
TypeError: Parameter to MergeFrom() must be instance of same class: expected tensorflow.TensorShapeProto got tensorflow.TensorShapeProto.
[11509] Failed to execute script 'main' due to unhandled exception: Parameter to MergeFrom() must be instance of same class: expected tensorflow.TensorShapeProto got tensorflow.TensorShapeProto.

成因分析 #

这个错误的核心是Python模块重复加载导致类型不匹配：

• TensorFlow依赖Protobuf库，PyInstaller打包时可能同时引入了系统环境的Protobuf和TensorFlow自带的Protobuf版本，或是打包过程中TensorFlow内部的Protobuf类被多次导入，导致内存中存在两个完全相同但实例独立的TensorShapeProto类，引发类型校验失败。
• MacOS的动态链接机制与PyInstaller打包逻辑交互时，更容易出现这类模块路径冲突，而Windows的打包机制对重复模块的处理更稳健。
• 加载外部模型时的AttributeError: as_proto错误，同样是Protobuf类导入异常导致无法调用序列化方法。

解决方法 #

方法1：使用--collect-all确保TensorFlow依赖完整收集 #

打包时添加--collect-all tensorflow参数，强制PyInstaller收集TensorFlow的所有子模块和依赖，避免模块缺失或重复加载：

pyinstaller --collect-all tensorflow --name tf_model_test your_script.py

方法2：修改PyInstaller Spec文件，指定Hidden Imports #

1. 先生成spec文件：

pyinstaller --name tf_model_test your_script.py

2. 打开生成的tf_model_test.spec文件，在hiddenimports列表中添加TensorFlow Protobuf相关模块：

hiddenimports=[
    'tensorflow.core.framework.tensor_shape_pb2',
    'tensorflow.core.framework.types_pb2',
    'tensorflow.core.framework.attr_value_pb2',
    'tensorflow.core.framework.graph_pb2',
    'tensorflow.core.framework.node_def_pb2',
    'tensorflow.core.framework.op_def_pb2',
    'tensorflow.core.framework.saved_model_pb2',
    'tensorflow.core.framework.version_pb2'
]

3. 基于修改后的spec文件打包：

pyinstaller tf_model_test.spec

方法3：使用--onedir模式打包 #

避免使用--onefile单文件模式，改用--onedir目录模式打包，减少临时解压时的模块路径冲突：

pyinstaller --onedir --collect-all tensorflow your_script.py

方法4：调整模型定义方式 #

将@dataclass装饰器与tf.Module的组合替换为纯tf.Module，避免打包时的序列化异常：

class ColorFilterModel(tf.Module):
    def __init__(self):
        self.rgb_filter_strengths = (0., 1., 1.)
    
    def compute_mahalonabis_sq_heatmap(self, image):
        return tf.reduce_sum(tf.cast(image, tf.float32) * tf.constant(self.rgb_filter_strengths), axis=-1)

验证方法 #

打包完成后，进入生成的目录（或运行单文件）执行测试脚本，确认模型保存和加载流程正常运行，无报错信息。

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