如何检测类中是否含TensorFlow对象并实现子类的高效序列化？

我之前也碰到过完全一样的场景——父类被多个子类继承，其中部分子类用到TensorFlow张量，直接用pickle序列化/反序列化踩了不少坑。下面分三个核心问题给你拆解解决方案，都是实际验证过的可行方案：

一、检查类是否包含TensorFlow对象 #

要判断一个类本身（不是实例）是否依赖TensorFlow对象，可以从三个维度入手：

• 检查类的属性/方法注解：如果有tf.Tensor或tf.Variable这类类型注解，直接识别
• 扫描类属性：如果类定义了静态TF张量/变量，直接判断
• 检查方法源码：借助inspect模块扫描类的方法源码，看是否引用了TensorFlow

给你一个实用的实现：

import tensorflow as tf
import inspect

def class_uses_tensorflow(cls):
    # 检查类的类型注解
    if hasattr(cls, '__annotations__'):
        for annotation in cls.__annotations__.values():
            # 处理普通注解和泛型注解（比如Optional[tf.Tensor]）
            if annotation is tf.Tensor or (hasattr(annotation, '__origin__') and annotation.__origin__ in (tf.Tensor, tf.Variable)):
                return True
    # 检查类的静态属性
    for attr_name, attr_val in cls.__dict__.items():
        if isinstance(attr_val, (tf.Tensor, tf.Variable)):
            return True
    # 扫描类方法的源码，判断是否引用TF
    for _, method in inspect.getmembers(cls, predicate=inspect.isfunction):
        try:
            source_code = inspect.getsource(method)
            if 'tf.' in source_code or 'tensorflow.' in source_code:
                return True
        except OSError:
            # 内置方法或无法获取源码的情况，直接跳过
            continue
    return False

二、检测对象是否包含TensorFlow属性 #

对于实例对象，需要递归遍历它的所有属性（包括私有属性），判断是否存在TF张量、变量或Module对象，同时要避免循环引用：

def obj_has_tensorflow_attrs(obj, visited=None):
    if visited is None:
        visited = set()
    # 记录对象ID，避免循环引用导致无限递归
    obj_id = id(obj)
    if obj_id in visited:
        return False
    visited.add(obj_id)
    
    # 直接判断当前对象是否是TF核心类型
    if isinstance(obj, (tf.Tensor, tf.Variable, tf.Module)):
        return True
    
    # 遍历所有属性，跳过内置特殊属性
    for attr_name in dir(obj):
        if attr_name.startswith('__') and attr_name.endswith('__'):
            continue
        try:
            attr_val = getattr(obj, attr_name)
            # 递归检查属性值
            if obj_has_tensorflow_attrs(attr_val, visited):
                return True
        except (AttributeError, TypeError):
            # 有些属性无法访问（比如@property的只读限制），直接跳过
            continue
    return False

三、高效保存/加载含TensorFlow对象的实例 #

直接用pickle处理TF对象容易报错，因为TF张量的序列化需要遵循自身的机制。这里推荐两种实用方案：

方案1：自定义Pickle钩子方法（适合非Module子类） #

在含TF对象的子类中实现__getstate__和__setstate__，把TF对象转换成numpy数组（可被pickle序列化），加载时再转回TF类型：

import pickle
import tensorflow as tf

class BaseClass:
    pass

class TFSubClass(BaseClass):
    def __init__(self, metadata, core_tensor):
        self.metadata = metadata
        self.core_tensor = core_tensor  # 可能是tf.Tensor或tf.Variable
    
    def __getstate__(self):
        # 复制对象状态，替换TF对象为可序列化格式
        state = self.__dict__.copy()
        if isinstance(state['core_tensor'], tf.Tensor):
            state['core_tensor'] = state['core_tensor'].numpy()
        elif isinstance(state['core_tensor'], tf.Variable):
            # 保存Variable的完整状态（数值、名称、可训练性）
            state['core_tensor'] = {
                'numpy_data': state['core_tensor'].numpy(),
                'name': state['core_tensor'].name,
                'trainable': state['core_tensor'].trainable
            }
        return state
    
    def __setstate__(self, state):
        # 恢复TF对象
        if isinstance(state['core_tensor'], dict):
            # 恢复tf.Variable
            state['core_tensor'] = tf.Variable(
                state['core_tensor']['numpy_data'],
                name=state['core_tensor']['name'],
                trainable=state['core_tensor']['trainable']
            )
        else:
            # 恢复tf.Tensor
            state['core_tensor'] = tf.convert_to_tensor(state['core_tensor'])
        self.__dict__.update(state)

# 测试代码
test_obj = TFSubClass("sample data", tf.Variable([1,2,3], trainable=True))
with open("tf_instance.pkl", "wb") as f:
    pickle.dump(test_obj, f)

with open("tf_instance.pkl", "rb") as f:
    loaded_obj = pickle.load(f)

print(loaded_obj.metadata)
print(loaded_obj.core_tensor)

方案2：结合TF SavedModel与Pickle（适合tf.Module子类） #

如果子类继承了tf.Module，推荐用tf.saved_model保存TF相关的权重和结构，用pickle保存非TF属性，这样能最大化兼容TF的生态：

import pickle
import tensorflow as tf
import os

class TFModuleSubClass(BaseClass, tf.Module):
    def __init__(self, metadata, trainable_var):
        super().__init__()
        self.metadata = metadata
        self.trainable_var = trainable_var
    
    def save(self, save_root_dir):
        # 创建保存目录
        os.makedirs(save_root_dir, exist_ok=True)
        # 保存TF模块部分
        tf.saved_model.save(self, os.path.join(save_root_dir, "tf_component"))
        # 保存非TF属性（过滤掉TF相关对象）
        non_tf_state = {
            k: v for k, v in self.__dict__.items() 
            if not isinstance(v, (tf.Tensor, tf.Variable, tf.Module))
        }
        with open(os.path.join(save_root_dir, "non_tf_data.pkl"), "wb") as f:
            pickle.dump(non_tf_state, f)
    
    @classmethod
    def load(cls, save_root_dir):
        # 加载TF模块
        tf_component = tf.saved_model.load(os.path.join(save_root_dir, "tf_component"))
        # 加载非TF属性
        with open(os.path.join(save_root_dir, "non_tf_data.pkl"), "rb") as f:
            non_tf_state = pickle.load(f)
        # 合并状态创建实例
        instance = cls.__new__(cls)
        instance.__dict__.update(tf_component.__dict__)
        instance.__dict__.update(non_tf_state)
        return instance

# 测试代码
test_obj = TFModuleSubClass("module sample", tf.Variable([4,5,6], trainable=True))
save_dir = "./tf_module_save"
test_obj.save(save_dir)
loaded_obj = TFModuleSubClass.load(save_dir)

print(loaded_obj.metadata)
print(loaded_obj.trainable_var)

额外优化：统一父类的保存逻辑 #

如果子类数量多，可以在父类中实现通用的保存方法，自动判断实例是否含TF属性，选择对应方案：

class BaseClass:
    def save(self, path):
        if obj_has_tensorflow_attrs(self):
            if isinstance(self, tf.Module):
                # 用方案2保存Module子类
                self.save(path)  # 这里调用子类的save方法
            else:
                # 用方案1的pickle保存
                with open(path, "wb") as f:
                    pickle.dump(self, f)
        else:
            # 普通子类直接用pickle
            with open(path, "wb") as f:
                pickle.dump(self, f)

这样所有子类都可以统一调用save方法，不需要每个子类单独写逻辑，维护起来更方便。

内容的提问来源于stack exchange，提问作者Sadjad Anzabi Zadeh