AssertionError: Tried to export a function which references untracked resource

Question

I wrote a unit-test in order to safe a model after noticing that I am not able to do so (anymore) during training.

@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)

Trying to save the model will always throw the following error:

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>

Note: As you can see in the code above, but I am not able to retrieve this tensor with tf.compat.v1.get_default_graph().get_tensor_by_name("77040:0").

I tried the following too, but the result is always empty:

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)

The problem is that I do not understand why I am getting this because the affected layer is tracked by Keras as you can see in the screenshot below. I took it during a debug-session in the call() function.

I have no explanation for this and I am running out of ideas what the issue might be here.

The transformations list in the screenshot is a property of and getting constructed by a layer InputsEmbedding like so:

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}")

---

In order to verify that the problem is not the InputsEmbedding, I created a unit-text for saving a model that is using just this particular layer.

@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)

Here I am able to save this layer without any issues:

ConvolutionStack

As it seems to be relevant, here is the (rather ugly) implementation of ConvolutionStack:

from typing import List

import tensorflow as tf
from tensorflow.keras import layers
from tensorflow.python.keras.layers import convolutional

from speech.lab.layers import InputsRequirements
from speech.lab.models import conv_util, models_util


class ConvolutionStack(layers.Layer):
    def __init__(
        self,
        convolutions: List[str],
        kernel_regularizer: dict = None,
        bias_regularizer: dict = None,
        **kwargs
    ):
        super().__init__(**kwargs)
        self.config = dict(
            convolutions=convolutions,
            kernel_regularizer=kernel_regularizer,
            bias_regularizer=bias_regularizer
        )
        self.conv_stack_config = [eval(f"dict({convolution})") for convolution in convolutions]
        self.conv_blocks = list()

        if kernel_regularizer is not None:
            kernel_regularizer = models_util.maybe_to_regularizer(kernel_regularizer)
        if bias_regularizer is not None:
            bias_regularizer = models_util.maybe_to_regularizer(bias_regularizer)

        for block_config in self.conv_stack_config:
            block = _new_convolution_block(
                **block_config,
                kernel_regularizer=kernel_regularizer,
                bias_regularizer=bias_regularizer,
            )
            self.conv_blocks.append(block)

        self.drop_dim2 = layers.Lambda(tf.squeeze, arguments=dict(axis=-2))
        self.expand_last = layers.Lambda(tf.expand_dims, arguments=dict(axis=-1))

    @property
    def inputs_requirements(self) -> InputsRequirements:
        requirements, frame_look_back = conv_util.get_conv2d_stack_requirements(self.conv_stack_config)
        first = requirements[0]
        t_min, f_size = first["min_size"]
        t_grow, f_grow = first["grow_size"]
        return InputsRequirements(
            frame_look_back=frame_look_back,
            t_min=t_min,
            t_grow=t_grow,
            f_min=f_size,
            f_grow=f_grow,
        )

    def call(self, inputs, training=None, mask=None, **kwargs):
        """
        :param inputs:
            Tensor taking the form [batch, time, freq, channel]
        :param training:
        :param mask:
        :param kwargs:
        :return:
            Tensor taking the form [batch, time, freq, 1]
        """

        if training:
            t_min = self.inputs_requirements.t_min
            t_grow = self.inputs_requirements.t_grow
            pad = conv_util.get_padding_for_loss(tf.shape(inputs)[1], t_min=t_min, t_grow=t_grow)
            inputs = tf.pad(inputs, ((0, 0), (0, pad), (0, 0), (0, 0)))

            if mask is not None:
                mask = tf.pad(mask, ((0, 0), (0, pad)))

        f_min = self.inputs_requirements.f_min
        f_grow = self.inputs_requirements.f_grow
        assert (inputs.shape[2] - f_min) % f_grow == 0, (
            f'Inputs dimension "freq" ' f"expected to be {f_min} + n * {f_grow}  but got {inputs.shape[2]} instead."
        )

        x = inputs
        for block in self.conv_blocks:

            for layer in block:

                if mask is not None and isinstance(layer, convolutional.Conv):
                    st, _ = layer.strides
                    kt = tf.maximum(layer.kernel_size[0] - 1, 1)
                    mask = mask[:, :-kt][:, ::st]
                    mask = tf.pad(mask, ((0, 0), (0, tf.maximum(2 - layer.kernel_size[0], 0))))

                x = layer(x, training=training)

        return self.expand_last(self.drop_dim2(x)), mask

    def get_config(self):
        return self.config


def _new_convolution_block(
    filters: int,
    kernel_size: tuple,
    strides: tuple,
    use_bias: bool = False,
    use_norm: bool = True,
    kernel_regularizer=None,
    bias_regularizer=None,
    activation=None,
):
    assert strides[0] % 2 == 0 or strides[0] == 1, "Strides on the time axis must be divisible by 2 or be exactly 1."

    if activation is not None:
        activation_layer = layers.Activation(activation)
    else:
        activation_layer = layers.Lambda(lambda x: x)

    if use_norm:
        norm_layer = layers.LayerNormalization()
    else:
        norm_layer = layers.Lambda(lambda x: x)

    return (
        layers.Conv2D(
            filters=filters,
            kernel_size=kernel_size,
            strides=strides,
            use_bias=use_bias,
            kernel_regularizer=kernel_regularizer,
            bias_regularizer=bias_regularizer,
        ),
        norm_layer,
        activation_layer,
    )

See also

• tensorflow/serving #1719

Pedant

Answer 1

Using

• tensorflow v2.5.0
• Python: 3.9

It appears that the problem occurs when we declare/define a layer as class-variable. I can only assume that the problem has to do with the internal Keras logic, which probably makes sense, but imo it's not obvious to the user and I don't think I have ever seen a hint pointing out that this can be an issue.

So, in my project I am having the following:

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 Example

The following example creates instances of ModelA, ModelB, ModelC and ModelD. Model A and B can be saved but C cannot. From what I can tell, it does not work to declarea layer which has trainable weights as class-variable. However, it does seem to work for layers which do not have trainable weights (see ModelB).

Please note how ModelD can be saved though. The difference to ModelB is that the layer gets only declared and not defined as None which leads to the question why ModelC works though.

Source Code

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()

Output

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!

Pedant

Thank you very much for providing such a clean code to reproduce the issue and understand its behavior – Mccullers

No problem. I hope it helps the next poor guy who runs into this issue. It was quite well hidden an the error message was not helpful at all. ^^ – Pedant

I think this is because keras layer does not keep track of class property, but it does track instance property, so when I added an useless self.tmp = self.model_layer in __init__, we can save it without problem. – Angio

I think it makes sense to not have a reference as a class property. Also in your example fixes, you don't really have one either, since with self.model_layer = ... you're not overwriting the class property, but just setting an instance property. I'm having this issue myself, without having any class property defined... – Slovak

Answer 2

Your issue is not related to 'transformer_transducer/transducer_encoder/inputs_embedding/ convolution_stack/conv2d/kernel:0'.
The error code tells you this element is referring to a non trackable element. It seems the non-trackable object is not directly assigned to an attribute of this conv2d/kernel:0.

To solve your issue, we need to localize Tensor("77040:0", shape=(), dtype=resource) from this error code:

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.

Edit:

Thanks to your comments, we found that "ConvolutionStack" seems to reproduce the error.

The problem only occurs if I use the ConvolutionStack layer in InputsEmbedding but I can save both of them successfully in a standalone model.

I understand you cannot share the config of this layer and that's why I suggest you try and localize this Tensor(77040:0 from the ConvolutionStack.

This untrackable tensor must be an artifcat or a temporary tensor created by a process of a function of ConvolutionStack.

Try to find a tensor that could be passed from a function to another instead of being assigned to an attribute of a layer's class

Mccullers

I feel like a noob asking this but how can I localize that tensor? – Pedant

@StefanFalk Could you provide the config of the layer conv2d/kernel:0 – Mccullers

conv2d is just a tf.keras.layers.Conv2D-layer. – Pedant

Take a look at the unit-test test_inputs_embedding_save_model(). All InputsEmbedding does is apply these convolutions by using the Keras Conv2D layer. And, as seen in the test, I am able to save a model which uses the InputsEmbedding layer. However, it's not working in with the actual model which I am trying to save and there's no indication where this tensor 77040:0 actually comes from. – Pedant

I'm also trying to understand where does this 77040:0 comes from. How is implemented the input layer in the actual implementation? Is it exactly the same way as your test inputs = layers.Input(shape=(None, num_features, num_channels)) ? – Mccullers

Yes. It's exactly the same. I wrote this unit-test just to compare and see whether or not this would crash too. To my surprise it's working. InputsEmbedding can do other things too e.g. StackFrames - this layer does not make any problems as well. The problem only occurs if I use the ConvolutionStack layer in InputsEmbedding but I can save both of them successfully in a standalone model. – Pedant

Ok very interesting. Could you provide the config that actually reproduce the issue ? – Mccullers

It's complicated. This layer is just the first input-layer of a Transducer model which I am using for speech recognition. Unfortunately I cannot share this code. But it's so weird to see that everything seems to work fine except in this particular case where I need it to work.. – Pedant

I've just added the code of ConvolutionStack - I must be overlooking something but to me it looks like there are no temporary tensors or anything. I am going to debug this today again, your guess is probably right, but so far I just didn't see it. At the beginning I thought the problem might be return value of ConvolutionStack. _new_convolution_block() but all these tensors are tracked by Keras so this can't be the problem (also, the unit-test for saving that layer is working). And if the problematic tensor comes from outside, why can I safe the model if I remove the ConvolutionStack? – Pedant

After a wild guess I was able to fix the problem (see). Thank you for helping me. I am going to accept my answer but I'll award you the bounty for helping me out :) – Pedant

Answer 3

Another possible solution is to remove slots from the model

class ModelA(tf.keras.Model):
    # this model raises the error in the question
    __slots__ = ["model_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()


class ModelB(tf.keras.Model):
    # this model does NOT raise the error in the question
    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()

The same goes with custom layers in tensorflow 2.9.1 and python 3.7

Trawick

Answer 4

My case for being stuck with a Tried to export a function which references untracked resource assertion error was related to using Lambda layer to transform Inputs prior to sending them to TextVectorization, while trying to use the train/inference model pattern.

Concretely: My model uses TextVectorization layer, but the inputs are numeric (int64 IDs). Because the inputs are not strings, I need to pass them through tf.strings.as_string(inputs) prior to calling the vectorizer with those inputs. Long story short, this is easily solved using the Lambda layer. The issue is that when you export the inference model with the Lambda fronting the TextVectorization layer, you can no longer track the TextVectorization layer, which in turn, needs to store the vocabulary data in a Variable.

My solution: create a custom layer that wraps the TextVectorization layer (as is actually recommended in the docs)

class NumericVectorization(tf.keras.layers.Layer):
    def __init__(self, vocab, output_sequence_length=1, name=None):
        super(NumericVectorization, self).__init__(name=name)
        self.vec = TextVectorization(
            standardize=None,
            split=None,
            output_mode='int',
            output_sequence_length=output_sequence_length,
            vocabulary=vocab
        )

    def call(self, inputs, *args, **kwargs):
        as_str = tf.strings.as_string(inputs)
        return self.vec(as_str)

Allpurpose

Answer 5

I faced similar issues with saving the model, however none of the solutions helped me directly.

With tensorflow 1.13.1 the following imports are fine, and the result could be saved with model.save(path):

from tensorflow.keras.layers import Input, Concatenate, ConvLSTM2D
from tensorflow.keras.layers import BatchNormalization
from tensorflow.python.keras.layers.convolutional import Conv2D
tensorflow.keras.models import Model

However, when I used this old project with newer tensorflow > 2.12 I got the following error:

AssertionError: Tried to export a function which references an 'untracked' 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.

Turned out, that in this newer version the import line:

from tensorflow.python.keras.layers.convolutional import Conv2D

Should look like this instead:

from tensorflow.keras.layers import Conv2D

This is because of the newer version of keras is available and imports are changed. Make sure to use tensorlow.keras... in similar cases.

Similar issue: https://mcmap.net/q/600519/-keras-cannot-import-name-conv2d

Guncotton

Answer 6

make sure to use all the packages with tensorflow.keras not keras

Bourgeoisie

Your answer could be improved with additional supporting information. Please edit to add further details, such as citations or documentation, so that others can confirm that your answer is correct. You can find more information on how to write good answers in the help center. – Sizable