Fix tf memory leak: tensorflow/tensorflow#37653 (comment)

Author: hankcs

hanlp/common/keras_component.py

@@ -11,7 +11,7 @@
 import tensorflow as tf
 
 import hanlp.utils
-from hanlp_common.io import save_json,load_json
+from hanlp_common.io import save_json, load_json
 from hanlp.callbacks.fine_csv_logger import FineCSVLogger
 from hanlp.common.component import Component
 from hanlp.common.transform_tf import Transform
@@ -255,7 +255,8 @@ def build_optimizer(self, optimizer, **kwargs):
         if isinstance(optimizer, (str, dict)):
             custom_objects = {'AdamWeightDecay': AdamWeightDecay}
             optimizer: tf.keras.optimizers.Optimizer = tf.keras.utils.deserialize_keras_object(optimizer,
-                                                                                               module_objects=vars(tf.keras.optimizers),
+                                                                                               module_objects=vars(
+                                                                                                   tf.keras.optimizers),
                                                                                                custom_objects=custom_objects)
         self.config.optimizer = tf.keras.utils.serialize_keras_object(optimizer)
         return optimizer
@@ -437,6 +438,7 @@ def predict(self, data: Any, batch_size=None, **kwargs):
             for output in self.predict_batch(batch, inputs=inputs, **kwargs):
                 results.append(output)
             num_samples += samples_in_batch
+        self.transform.cleanup()
 
         if flat:
             return results[0]

hanlp/common/transform_tf.py

@@ -28,6 +28,8 @@ def __init__(self, config: SerializableDict = None, map_x=True, map_y=True, **kw
         self.output_types = None
         self.output_shapes = None
         self.padding_values = None
+        # Fix tf memory leak: https://github.com/tensorflow/tensorflow/issues/37653#issuecomment-1000517720
+        self.py_func_set_to_cleanup = set()
 
     @abstractmethod
     def fit(self, trn_path: str, **kwargs) -> int:
@@ -170,6 +172,9 @@ def samples_to_dataset(self, samples: Generator, map_x=None, map_y=None, batch_s
                                padding_values]), 'Your create_types_shapes_values returns None, which is not allowed'
         # if not callable(samples):
         #     samples = Transform.generator_to_callable(samples)
+        if not hasattr(tf.compat.v1.get_default_graph(), '_py_funcs_used_in_graph'):
+            tf.compat.v1.get_default_graph()._py_funcs_used_in_graph = []
+        py_func_set_before = set(tf.compat.v1.get_default_graph()._py_funcs_used_in_graph)
         dataset = tf.data.Dataset.from_generator(samples, output_types=output_types, output_shapes=output_shapes)
         if cache:
             logger.debug('Dataset cache enabled')
@@ -197,6 +202,8 @@ def mapper(X, Y):
                 return X, Y
 
             dataset = dataset.map(mapper, num_parallel_calls=tf.data.experimental.AUTOTUNE)
+        py_func_set_after = set(tf.compat.v1.get_default_graph()._py_funcs_used_in_graph) - py_func_set_before
+        self.py_func_set_to_cleanup |= py_func_set_after
         return dataset
 
     @abstractmethod
@@ -237,7 +244,8 @@ def str_to_idx(self, X, Y) -> Tuple[Union[tf.Tensor, Tuple], tf.Tensor]:
     def X_to_inputs(self, X: Union[tf.Tensor, Tuple[tf.Tensor]]) -> Iterable:
         return [repr(x) for x in X]
 
-    def Y_to_outputs(self, Y: Union[tf.Tensor, Tuple[tf.Tensor]], gold=False, inputs=None, X=None, batch=None) -> Iterable:
+    def Y_to_outputs(self, Y: Union[tf.Tensor, Tuple[tf.Tensor]], gold=False, inputs=None, X=None,
+                     batch=None) -> Iterable:
         return [repr(y) for y in Y]
 
     def XY_to_inputs_outputs(self, X: Union[tf.Tensor, Tuple[tf.Tensor]],
@@ -295,3 +303,8 @@ def input_truth_output_to_str(self, input, truth, output):
 
         """
         return '\t'.join([input, truth, output]) + '\n'
+
+    def cleanup(self):
+        new_py_funcs = set(tf.compat.v1.get_default_graph()._py_funcs_used_in_graph) - self.py_func_set_to_cleanup
+        tf.compat.v1.get_default_graph()._py_funcs_used_in_graph = list(new_py_funcs)
+        self.py_func_set_to_cleanup = set()