tensorflow
when tf.keras meets get_model_type a memory leakage happens
System information
- Have I specified the code to reproduce the issue (Yes/No): yes
- Environment in which the code is executed (e.g., Local (Linux/MacOS/Windows), Interactive Notebook, Google Cloud, etc): - TensorFlow version (you are using): 2.3.2- TFX Version: 0.26.1- Python version:3.6.7
Describe the current behavior In tfma module, model_util.py :109 get_model_type function, it will load model into gpu and inference the type of model then return model's type.
if model_path:
try:
keras_model = tf.keras.models.load_model(model_path)
# In some cases, tf.keras.models.load_model can successfully load a
# saved_model but it won't actually be a keras model.
if isinstance(keras_model, tf.keras.models.Model):
return constants.TF_KERAS
except Exception: # pylint: disable=broad-except
pass
if tags:
if tags and eval_constants.EVAL_TAG in tags:
return constants.TF_ESTIMATOR
else:
return constants.TF_GENERIC
signature_name = None
if model_spec:
if model_spec.signature_name:
signature_name = model_spec.signature_name
else:
signature_name = tf.saved_model.DEFAULT_SERVING_SIGNATURE_DEF_KEY
if signature_name == eval_constants.EVAL_TAG:
return constants.TF_ESTIMATOR
else:
return constants.TF_GENERIC
But thanks to this function, the loaded model keep staying in the gpu memory and takes large amount of memory,
when the evaluator goes on, it starts to do the evaluation, another load process is started and cause an OOM error
Describe the expected behavior When the get model type finished the GPU memory should be released or we just keep passing the loaded model to the following process, instead load it again.
Standalone code to reproduce the issue Providing a bare minimum test case or step(s) to reproduce the problem will greatly help us to debug the issue. If possible, please share a link to Colab/Jupyter/any notebook. Just try a larger pretrained model like multilingual bert and put it into pipeline with evaluator.
Name of your Organization (Optional)
Other info / logs Include any logs or source code that would be helpful to
diagnose the problem. If including tracebacks, please include the full
traceback. Large logs and files should be attached.
""
RuntimeError: Traceback (most recent call last):
File "apache_beam/runners/common.py", line 1280, in apache_beam.runners.common.DoFnRunner._invoke_lifecycle_method
File "apache_beam/runners/common.py", line 516, in apache_beam.runners.common.DoFnInvoker.invoke_setup
File "/home/luban/.local/lib/python3.6/site-packages/tensorflow_model_analysis/model_util.py", line 758, in setup
super(ModelSignaturesDoFn, self).setup()
File "/home/luban/.local/lib/python3.6/site-packages/tensorflow_model_analysis/model_util.py", line 574, in setup
model_load_time_callback=self._set_model_load_seconds)
File "/home/luban/.local/lib/python3.6/site-packages/tensorflow_model_analysis/types.py", line 160, in load
return self._shared_handle.acquire(construct_fn)
File "/home/luban/.local/lib/python3.6/site-packages/tfx_bsl/beam/shared.py", line 238, in acquire
return _shared_map.acquire(self._key, constructor_fn, tag)
File "/home/luban/.local/lib/python3.6/site-packages/tfx_bsl/beam/shared.py", line 194, in acquire
result = control_block.acquire(constructor_fn, tag)
File "/home/luban/.local/lib/python3.6/site-packages/tfx_bsl/beam/shared.py", line 89, in acquire
result = constructor_fn()
File "/home/luban/.local/lib/python3.6/site-packages/tensorflow_model_analysis/types.py", line 169, in with_load_times
model = self.construct_fn()
File "/home/luban/.local/lib/python3.6/site-packages/tensorflow_model_analysis/model_util.py", line 550, in construct_fn
model = tf.compat.v1.saved_model.load_v2(eval_saved_model_path, tags=tags)
File "/home/luban/.local/lib/python3.6/site-packages/tensorflow/python/saved_model/load.py", line 603, in load
return load_internal(export_dir, tags, options)
File "/home/luban/.local/lib/python3.6/site-packages/tensorflow/python/saved_model/load.py", line 633, in load_internal
ckpt_options)
File "/home/luban/.local/lib/python3.6/site-packages/tensorflow/python/saved_model/load.py", line 130, in init
self._load_all()
File "/home/luban/.local/lib/python3.6/site-packages/tensorflow/python/saved_model/load.py", line 141, in _load_all
self._load_nodes()
File "/home/luban/.local/lib/python3.6/site-packages/tensorflow/python/saved_model/load.py", line 296, in _load_nodes
slot_name=slot_variable_proto.slot_name)
File "/home/luban/.local/lib/python3.6/site-packages/tensorflow/python/keras/optimizer_v2/optimizer_v2.py", line 764, in add_slot
initial_value=initial_value)
File "/home/luban/.local/lib/python3.6/site-packages/tensorflow/python/ops/variables.py", line 262, in call
return cls._variable_v2_call(*args, **kwargs)
File "/home/luban/.local/lib/python3.6/site-packages/tensorflow/python/ops/variables.py", line 256, in _variable_v2_call
shape=shape)
File "/home/luban/.local/lib/python3.6/site-packages/tensorflow/python/ops/variables.py", line 237, in
i just suggest that all the functions not the only one I mentioned above but all the functions including tfma/api/model_eval_lib.py：380 should fix this once for all
I'm not sure how to control the memory in this case. The model is loaded and then thrown away which should release the memory. This is likely related to [1]. For now you can set the model_type explicitly in the ModelSpec to avoid doing any work in the get_model_type call.
[1] https://github.com/tensorflow/tensorflow/issues/36465
I'm not sure how to control the memory in this case. The model is loaded and then thrown away which should release the memory. This is likely related to [1]. For now you can set the model_type explicitly in the ModelSpec to avoid doing any work in the get_model_type call.
yep , still as it said in the quote, looks like this clumsy way that tf manage gpu memory keeps being a needle under the bed :(