StringLookup layer never releases memory on load about tf-keras

[cjmcgraw]

One potential solution I've found is to manage the TrackableResources directly myself.

Consider the following code:

prev_model = loaded_model
loaded_model = tf.saved_model.load("/tmp/test-model")

if prev_model:
    lookups = filter(None, [
        layer._trackable_children().get('lookup_table') if hasattr(layer, '_trackable_children') else None
        for layer in prev_model._trackable_children().values()
    ])
    for lookup in lookups:
        tf.raw_ops.DestroyResourceOp(resource=lookup.resource_handle)

    del prev_model

Doing that after reloading appears to clear the resources as expected and not cause memory to grow indefinitely.

[cjmcgraw]

For any future readers. V2:

def wipeit(model):
    def _give_handles(m):
        if hasattr(m, 'resource_handle'):
            yield m.resource_handle
    
        children =  {}
        try:
            children = m._trackable_children()
        except:
            pass

        for child in children.values():
            yield from _give_handles(child)


    for handle in _give_handles(model):
        tf.raw_ops.DestroyResourceOp(resource=handle)


...

# during model load, after its loaded
prev_model = loaded_model
loaded_model = tf.saved_model.load(...)

if prev_model:
    wipeit(prev_model)
    del prev_model

[divyashreepathihalli]

@cjmcgraw thanks for filing the issue. This issues seems to be related to TF. This behavior can be reproduced without using any keras APIs. You can just use a lookup table and save and load the model -- the behavior would be the same.