AssertionError：尝试导出引用未跟踪资源的函数

问题描述

在注意到我在培训期间不能(再)这样做后，我编写了一个单元测试来保护模型。

@pytest.mark.usefixtures("maybe_run_functions_eagerly")
def test_save_model(speech_model: Tuple[TransducerBase, SpeechFeaturesConfig]):
    model, speech_features_config = speech_model
    speech_features_config: SpeechFeaturesConfig
    channels = 3 if speech_features_config.add_delta_deltas else 1
    num_mel_bins = speech_features_config.num_mel_bins
    enc_inputs = np.random.rand(1, 50, num_mel_bins, channels)
    dec_inputs = np.expand_dims(np.random.randint(0, 25, size=10), axis=1)
    inputs = enc_inputs, dec_inputs
    model(inputs)

    # Throws KeyError:
    # graph = tf.compat.v1.get_default_graph()
    # tensor = graph.get_tensor_by_name("77040:0")

    directory = tempfile.mkdtemp(prefix=f"{model.__class__.__name__}_")
    try:
        model.save(directory)
    finally:
        shutil.rmtree(directory)

尝试保存模型始终会引发以下错误：

E         AssertionError: Tried to export a function which references untracked resource Tensor("77040:0", shape=(), dtype=resource). TensorFlow objects (e.g. tf.Variable) captured by functions must be tracked by assigning them to an attribute of a tracked object or assigned to an attribute of the main object directly.
E         
E         Trackable Python objects referring to this tensor (from gc.get_referrers, limited to two hops):
E         <tf.Variable 'transformer_transducer/transducer_encoder/inputs_embedding/convolution_stack/conv2d/kernel:0' shape=(3, 3, 3, 32) dtype=float32>

注意：如您在上面的代码中所见，但我无法使用tf.compat.v1.get_default_graph().get_tensor_by_name("77040:0")检索此张量。

我也尝试了以下方法，但结果始终为空：

model(batch)  # Build the model

tensor_name = "77040"

var_names = [var.name for var in model.trainable_weights]
weights = list(filter(lambda var: tensor_name in var, var_names))

var_names = [var.name for var in model.trainable_variables]
variables = list(filter(lambda var: tensor_name in var, var_names))

print(weights)
print(variables)

问题是我不明白为什么会出现这种情况，因为受影响的层是由Kera跟踪的，正如您在下面的屏幕截图中所看到的那样。我是在call()函数的调试会话中使用它的。

我对此无法解释，我也想不出这里可能出了什么问题。

截图中的transformations列表是层的属性和构造InputsEmbedding，如下所示：

class InputsEmbedding(layers.Layer, TimeReduction):
    def __init__(self, config: InputsEmbeddingConfig, **kwargs):
        super().__init__(**kwargs)

        if config.transformations is None or not len(config.transformations):
            raise RuntimeError("No transformations provided.")

        self.config = config

        self.transformations = list()
        for transformation in self.config.transformations:
            layer_name, layer_params = list(transformation.items())[0]
            layer = _get_layer(layer_name, layer_params)
            self.transformations.append(layer)

        self.init_time_reduction_layer()

    def get_config(self):
        return self.config.dict()


def _get_layer(name: str, params: dict) -> layers.Layer:
    if name == "conv2d_stack":
        return ConvolutionStack(**params)
    elif name == "stack_frames":
        return StackFrames(**params)
    else:
        raise RuntimeError(f"Unsupported or unknown time-reduction layer {name}")

为了验证问题不是InputsEmbedding，我创建了一个单元文本来保存仅使用该特定层的模型。

@pytest.mark.usefixtures("maybe_run_functions_eagerly")
def test_inputs_embedding_save_model():
    convolutions = [
        "filters=2, kernel_size=(3, 3), strides=(2, 1)",
        "filters=4, kernel_size=(3, 3), strides=(2, 1)",
        "filters=8, kernel_size=(3, 4), strides=(1, 1)",
    ]

    config = InputsEmbeddingConfig()
    config.transformations = [dict(conv2d_stack=dict(convolutions=convolutions)), dict(stack_frames=dict(n=2))]

    num_features = 8
    num_channels = 3

    inputs = layers.Input(shape=(None, num_features, num_channels))
    x = inputs
    x, _ = InputsEmbedding(config)(x)
    model = keras.Model(inputs=inputs, outputs=x)
    model.build(input_shape=(1, 20, num_features, num_channels))

    directory = tempfile.mkdtemp(prefix=f"{model.__class__.__name__}_")
    try:
        model.save(directory)
    finally:
        shutil.rmtree(directory)

在这里，我可以毫无问题地保存该层：

另请参阅

• tensorflow/serving #1719

推荐答案

使用

• TensorFlow v2.5.0
• Python：3.9

问题似乎发生在我们将层声明/定义为类变量时。我只能假设问题与KERAS的内部逻辑有关，这可能是有道理的，但对于用户来说，这并不明显，我认为我从未看到任何提示指出这可能是一个问题。

因此，在我的项目中，我拥有以下内容：

class Model(keras.Model):
    inputs_embedding: InputsEmbedding = None  # <-- This caused the problem

    def __init__(config, *args, **kwargs):
        super().__init__(*args, **kwargs)
        if config.embeddings is not None:
            self.inputs_embedding = InputsEmbedding(config.embeddings)
        # ...

MVP示例

import tempfile

import numpy as np
import tensorflow as tf
from tensorflow.keras import layers


class ModelA(tf.keras.Model):

    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)
        self.model_layer = layers.LayerNormalization()

    def call(self, inputs, training=None, mask=None):
        return self.model_layer(inputs)

    def get_config(self):
        return dict()


class ModelB(tf.keras.Model):
    model_layer: layers.Layer = None

    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)
        # This is probably working because layers.Lambda has no trainable variables
        self.model_layer = layers.Lambda(lambda x: x)

    def call(self, inputs, training=None, mask=None):
        return self.model_layer(inputs)

    def get_config(self):
        return dict()


class ModelC(tf.keras.Model):
    model_layer: layers.Layer = None

    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)
        self.model_layer = layers.LayerNormalization()

    def call(self, inputs, training=None, mask=None):
        return self.model_layer(inputs)

    def get_config(self):
        return dict()


class ModelD(tf.keras.Model):
    model_layer: layers.Layer

    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)
        self.model_layer = layers.LayerNormalization()

    def call(self, inputs, training=None, mask=None):
        return self.model_layer(inputs)

    def get_config(self):
        return dict()


def save_tmp_model(model: tf.keras.Model):
    name = model.__class__.__name__
    print(f'Saving model {name}')
    try:
        model.save(tempfile.mkdtemp(prefix=f"{name}_"))
    except Exception as e:
        print(f"Unable to save model: {name}")
        print('Error message:')
        print(str(e))
        return

    print(f".. success!")


def main():
    inputs = np.random.rand(1, 50, 16)

    model_a = ModelA()
    model_b = ModelB()
    model_c = ModelC()
    model_d = ModelD()

    # Build models
    model_a(inputs)
    model_b(inputs)
    model_c(inputs)
    model_d(inputs)

    # Save models
    save_tmp_model(model_a)
    save_tmp_model(model_b)
    save_tmp_model(model_c)
    save_tmp_model(model_d)


if __name__ == '__main__':
    main()

Saving model ModelA
.. success!
Saving model ModelB
.. success!
Saving model ModelC
Unable to save model: ModelC
Error message:
Tried to export a function which references untracked resource Tensor("1198:0", shape=(), dtype=resource). TensorFlow objects (e.g. tf.Variable) captured by functions must be tracked by assigning them to an attribute of a tracked object or assigned to an attribute of the main object directly.

Trackable Python objects referring to this tensor (from gc.get_referrers, limited to two hops):
<tf.Variable 'model_c/layer_normalization_1/gamma:0' shape=(16,) dtype=float32>
Saving model ModelD
.. success!