Skip to content

chore: cache deserialized_module and script_module #4196

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 5 commits into from
Oct 11, 2024
Merged

chore: cache deserialized_module and script_module #4196

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
5 commits
Select commit Hold shift + click to select a range
95e4adf
chore: cache `deserialized_module` and `script_module`
njzjz Oct 8, 2024
013883c
proxy
njzjz Oct 8, 2024
bf6db03
xxxxx
njzjz Oct 8, 2024
f270632
clear cache
njzjz Oct 8, 2024
fac7a26
try lru_cache(maxsize=1)
njzjz Oct 9, 2024
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
18 changes: 18 additions & 0 deletions source/tests/universal/dpmodel/backend.py
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,4 +1,8 @@
# SPDX-License-Identifier: LGPL-3.0-or-later
from functools import (
lru_cache,
)

import numpy as np

from deepmd.dpmodel.common import (
Expand Down Expand Up @@ -30,8 +34,15 @@ def convert_to_numpy(cls, xx: np.ndarray) -> np.ndarray:
def convert_from_numpy(cls, xx: np.ndarray) -> np.ndarray:
return xx

@classmethod
@lru_cache(maxsize=1)
def _get_deserialized_module(cls):
return cls.module.deserialize(cls.module.serialize())

@property
def deserialized_module(self):
if hasattr(self.__class__, "module"):
return self._get_deserialized_module()
return self.module.deserialize(self.module.serialize())

@property
Expand All @@ -41,3 +52,10 @@ def modules_to_test(self):
self.deserialized_module,
]
return modules

@classmethod
def tearDownClass(cls):
super().tearDownClass()
if hasattr(cls, "module"):
del cls.module
cls._get_deserialized_module.cache_clear()
27 changes: 27 additions & 0 deletions source/tests/universal/pt/backend.py
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,4 +1,8 @@
# SPDX-License-Identifier: LGPL-3.0-or-later
from functools import (
lru_cache,
)

import numpy as np
import torch

Expand All @@ -18,13 +22,28 @@ class PTTestCase(BackendTestCase):
module: "torch.nn.Module"
"""PT module to test."""

@classmethod
@lru_cache(maxsize=1)
def _get_script_module(cls):
with torch.jit.optimized_execution(False):
return torch.jit.script(cls.module)

@property
def script_module(self):
if hasattr(self.__class__, "module"):
return self._get_script_module()
with torch.jit.optimized_execution(False):
return torch.jit.script(self.module)

@classmethod
@lru_cache(maxsize=1)
def _get_deserialized_module(cls):
return cls.module.deserialize(cls.module.serialize())

@property
def deserialized_module(self):
if hasattr(self.__class__, "module"):
return self._get_deserialized_module()
return self.module.deserialize(self.module.serialize())

@property
Expand All @@ -35,6 +54,14 @@ def modules_to_test(self):
]
return modules

@classmethod
def tearDownClass(cls):
super().tearDownClass()
if hasattr(cls, "module"):
del cls.module
cls._get_deserialized_module.cache_clear()
cls._get_script_module.cache_clear()

def test_jit(self):
if getattr(self, "skip_test_jit", False):
self.skipTest("Skip test jit.")
Expand Down
Loading