Jit pre save hook

python/paddle/fluid/dygraph/jit.py

@@ -1,4 +1,5 @@
 # Copyright (c) 2019 PaddlePaddle Authors. All Rights Reserved.
+# Copyright (c) 2021 NVIDIA Corporation. All rights reserved.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.
@@ -20,6 +21,7 @@
 import functools
 from collections import OrderedDict
 import inspect
+import threading
 
 import six
 import paddle
@@ -42,7 +44,8 @@
 
 __all__ = [
     'TracedLayer', 'declarative', 'dygraph_to_static_func', 'set_code_level',
-    'set_verbosity', 'save', 'load', 'not_to_static'
+    'set_verbosity', 'save', 'load', 'not_to_static', 'register_save_pre_hook',
+    'clear_save_pre_hooks'
 ]
 
 
@@ -525,6 +528,58 @@ def _build_load_path_and_config(path, config):
     return model_path, config
 
 
+_save_pre_hooks_lock = threading.Lock()
+_save_pre_hooks = []
+
+
+class HookRemoveHelper(object):
+    """ A HookRemoveHelper that can be used to remove hook. """
+
+    def __init__(self, hook):
+        self._hook = hook
+
+    def remove(self):
+        remove_save_pre_hook(self._hook)
+
+
+def register_save_pre_hook(hook):
+    global _save_pre_hooks_lock
+    global _save_pre_hooks
+    _save_pre_hooks_lock.acquire()
+    if hook not in _save_pre_hooks:
+        _save_pre_hooks.append(hook)
+    _save_pre_hooks_lock.release()
+    return HookRemoveHelper(hook)
+
+
+def remove_save_pre_hook(hook):
+    global _save_pre_hooks_lock
+    global _save_pre_hooks
+    _save_pre_hooks_lock.acquire()
+    if hook in _save_pre_hooks:
+        _save_pre_hooks.remove(hook)
+    _save_pre_hooks_lock.release()
+
+
+def clear_save_pre_hooks():
+    global _save_pre_hooks_lock
+    global _save_pre_hooks
+    _save_pre_hooks_lock.acquire()
+    _save_pre_hooks.clear()
+    _save_pre_hooks_lock.release()
+
+
+def run_save_pre_hooks(func):
+    def wrapper(layer, path, input_spec=None, **configs):
+        global _save_pre_hooks
+        for hook in _save_pre_hooks:
+            hook(layer, input_spec, configs)
+        func(layer, path, input_spec, **configs)
+
+    return wrapper
+
+
+@run_save_pre_hooks
 @switch_to_static_graph
 def save(layer, path, input_spec=None, **configs):
     """

python/paddle/fluid/tests/unittests/test_jit_pre_save_hooks.py

@@ -0,0 +1,61 @@
+# Copyright (c) 2021 PaddlePaddle Authors. All Rights Reserved.
+# Copyright (c) 2021 NVIDIA Corporation. All rights reserved.
+# 
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+# 
+#     http://www.apache.org/licenses/LICENSE-2.0
+# 
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from __future__ import print_function
+
+import unittest
+
+import paddle
+from paddle.jit import register_save_pre_hook
+from paddle.jit import clear_save_pre_hooks
+from paddle.fluid.dygraph.jit import run_save_pre_hooks
+
+_counter = 0
+
+
+class TestPreSaveHooks(unittest.TestCase):
+    def test_pre_save_hook_functions(self):
+        def fake_func(*args, **kwgs):
+            global _counter
+            _counter += 1
+
+        remove_handler = register_save_pre_hook(fake_func)
+        self.assertEqual(len(paddle.fluid.dygraph.jit._save_pre_hooks), 1)
+        self.assertTrue(
+            paddle.fluid.dygraph.jit._save_pre_hooks[0] is fake_func)
+
+        # Test of avoiding redundancy hanging
+        remove_handler = register_save_pre_hook(fake_func)
+        self.assertEqual(len(paddle.fluid.dygraph.jit._save_pre_hooks), 1)
+        self.assertTrue(
+            paddle.fluid.dygraph.jit._save_pre_hooks[0] is fake_func)
+
+        remove_handler.remove()
+        self.assertEqual(len(paddle.fluid.dygraph.jit._save_pre_hooks), 0)
+
+        remove_handler = register_save_pre_hook(fake_func)
+        clear_save_pre_hooks()
+        self.assertEqual(len(paddle.fluid.dygraph.jit._save_pre_hooks), 0)
+
+        global _counter
+        _counter = 0
+        remove_handler = register_save_pre_hook(fake_func)
+        func_with_hook = run_save_pre_hooks(fake_func)
+        func_with_hook(None, None)
+        self.assertEqual(_counter, 2)
+
+
+if __name__ == '__main__':
+    unittest.main()

python/paddle/jit/__init__.py

@@ -1,4 +1,5 @@
-#   Copyright (c) 2020 PaddlePaddle Authors. All Rights Reserved.
+# Copyright (c) 2020 PaddlePaddle Authors. All Rights Reserved.
+# Copyright (c) 2021 NVIDIA Corporation. All rights reserved.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.
@@ -21,6 +22,8 @@
 from ..fluid.dygraph.jit import set_verbosity  # noqa: F401
 from ..fluid.dygraph.jit import declarative as to_static  # noqa: F401
 from ..fluid.dygraph.jit import not_to_static  # noqa: F401
+from ..fluid.dygraph.jit import register_save_pre_hook  # noqa: F401
+from ..fluid.dygraph.jit import clear_save_pre_hooks  # noqa: F401
 from ..fluid.dygraph import ProgramTranslator  # noqa: F401
 from ..fluid.dygraph.io import TranslatedLayer  # noqa: F401
 
@@ -35,5 +38,7 @@
     'TranslatedLayer',
     'set_code_level',
     'set_verbosity',
-    'not_to_static'
+    'not_to_static',
+    'register_save_pre_hook',
+    'clear_save_pre_hooks'
 ]