@spate141/ All,

AssertionError: Tried to export a function which references 'untracked' resource Tensor("308003:0", shape=(), dtype=resource). can be solved by not defining trainable layers as class attributes in your sub-classes as Keras doesn’t track class, but only tracks layer instance as per commit https://github.com/tensorflow/tensorflow/commit/9d724a8e6034d321e97cdc9972d4d6e7adb3e3ca. You can refer here for clear explanation.

Also, you can try saving static HashTable as one of model’s properties as shown in above comment.

Thank you!

I have tried to export my text-classification model, built and trained using tf.keras as shown below, and I get the same error.

I have used tensorflow 2.7.0 on Ubuntu 18.04(google colab) to train and save the model.

my code:

class TFModel(tf.Module):
    def __init__(self, model: tf.keras.Model) -> None:
        self.model = model

        # this "key_value_tensor_initializer" is an instance of 'tf.lookup.KeyValueTensorInitializer', but this key-value initialiser is not build inside this constructor, but an outside global variable
        self.hash_table = tf.lookup.StaticHashTable(initializer=key_value_tensor_initializer,
                                                    default_value=1, name='hash_table')

    def pre_process(self, comment):
        # this function does all the pre-processing on the input payload to the format acceptable by model for predictions
        # convert the text to lower using "tf.strings.lower" -> does 3 regex ops with "tf.strings.regex_replace" one after the other -> 
        # splits sentences to words using "tf.strings.split" -> then look-up those in hash-table -> then pad 0s at the end -> return it
        return preprocessed_comment
        
    def post_process(self, probabilities):
        # this function converts the model predictions(probabilities) to output classes
        # this function will return a RaggedTensor
        return final_labels

    def f(self):
        @tf.function
        def inner(comment: List[str]) -> Dict[str, Union[List[float], List[ByteString], str]]:
            processed_comment = self.pre_process(comment)
            predicted_probs = self.model(processed_comment)
            predicted_labels = self.post_process(predicted_probs)
            return {'predicted_probabilities': predicted_probs,
                    'predicted_labels':predicted_labels,
                    'description': 'prediction probabilities ranges from 0 (hesitant) to 1 (confident).'}
        return inner

    ########## also tried, but to no avail ##########
    # @tf.function(input_signature=[tf.TensorSpec(shape=[1,], dtype=tf.string)])
    # @tf.function
    # def serving_default(self, comment: List[str]) -> Dict[str, Union[List[float], List[ByteString], str]]:
    #     processed_comment = self.pre_process(comment)
    #     predicted_probs = self.model(processed_comment)
    #     predicted_labels = self.post_process(predicted_probs)
    #     return {'predicted_probabilities': predicted_probs,
    #             'predicted_labels':predicted_labels,
    #             'description': 'prediction probabilities ranges from 0 (hesitant) to 1 (confident).'}
    ########################################


    
##### serialising part

## create an instance
tf_model_wrapper = TFModel(model)
# trying to create concrete_function as mentioned on github issue
concrete_fn = tf_model_wrapper.f().get_concrete_function(comment=tf.TensorSpec([None], tf.string))
## save the model to disk(serialize it)
tf.keras.models.save_model(
    model=tf_model_wrapper.model,
    filepath='/content/complex_nw_v1',
    signatures={'serving_default': concrete_fn})

but, I get this error:

---------------------------------------------------------------------------
AssertionError                            Traceback (most recent call last)
<ipython-input-76-889bb0fb80ba> in <module>()
    101     model=tf_model_wrapper.model,
    102     filepath='/content/complex_nw_v1',
--> 103     signatures={"serving_default": concrete_fn})

1 frames
/usr/local/lib/python3.7/dist-packages/keras/utils/traceback_utils.py in error_handler(*args, **kwargs)
     65     except Exception as e:  # pylint: disable=broad-except
     66       filtered_tb = _process_traceback_frames(e.__traceback__)
---> 67       raise e.with_traceback(filtered_tb) from None
     68     finally:
     69       del filtered_tb

/usr/local/lib/python3.7/dist-packages/tensorflow/python/saved_model/save.py in _map_captures_to_created_tensors(original_captures, resource_map)
    530           "directly.\n\n Trackable Python objects referring to this tensor "
    531           "(from gc.get_referrers, limited to two hops):\n{}".format("\n".join(
--> 532               [repr(obj) for obj in trackable_referrers])))
    533     export_captures.append(mapped_resource)
    534   return export_captures

AssertionError: Tried to export a function which references 'untracked' resource Tensor("308003: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):
<tensorflow.python.ops.lookup_ops.StaticHashTable object at 0x7f71763e1550>

How do I fix this and save the model with properly?

@naveen-marthala

I may be able to offer a solution, as I have recently encountered this error in a similar way. It turns out we need to save the HashTable as one of model’s properties. In this specific example,

## create an instance
tf_model_wrapper = TFModel(model)
# trying to create concrete_function as mentioned on github issue
concrete_fn = tf_model_wrapper.f().get_concrete_function(comment=tf.TensorSpec([None], tf.string))
## save the model to disk(serialize it)
model_to_save = tf_model_wrapper.model
model_to_save.hash_table = tf_model_wrapper.hash_table
tf.keras.models.save_model(
    model=model_to_save,
    filepath='/content/complex_nw_v1',
    signatures={'serving_default': concrete_fn})

The actual name of the attribute probably doesn’t matter.

Seems like I solved it! I would appreciate if you can add a version of this to the documentation so other can use it. To make the variables and the other elements from outside trackable, we need to write the Lambda layer using the subclassed Layer: tf.keras.layers.Layer

class LabelConverter(tf.keras.layers.Layer):
    
    def __init__(self, **kwargs):
        super(LabelConverter, self).__init__(**kwargs)

        # Implement your StaticHashTable here
        keys = tf.constant([0, 1], dtype=tf.int32)
        values = tf.constant([1000, 2000], dtype=tf.float32)
        table_init = tf.lookup.KeyValueTensorInitializer(keys, values)
        self.table = tf.lookup.StaticHashTable(table_init, -1)
    
    def build(self, input_shape):
        self.built = True
    
    def call(self, tensor_input):
        # this block is doing the transformation on input
        label_tensor = tf.cast(tensor_input[:, 0], tf.int32)
        score_tensor = tensor_input[:, 1]
        categories_tensor = self.table.lookup(label_tensor)
        return tf.stack((categories_tensor, score_tensor), axis=1)

• Exporting TF2 model with StaticHashTable for Serving:

# adding on top of already trained keras model
extra_layer = LabelConverter()(model.output)
hash_table_model = tf.keras.models.Model(inputs=model.input, outputs=extra_layer)

version = 1
name = 'tmp_test_serving'
export_path = f'/data/{name}/{version}'
tf.saved_model.save(hash_table_model, export_path)

EDIT: Doesn’t work on TF-Serving for some reason!