MCPcopy Create free account
hub / github.com/OpenGVLab/HumanBench / AmpHandle

Class AmpHandle

PATH/core/fp16/amp.py:59–102  ·  view source on GitHub ↗

Source from the content-addressed store, hash-verified

57 _USER_FLOAT_MODULE.add(module)
58
59class AmpHandle(object):
60 def __init__(self, enable_caching=True, verbose=False):
61 self._enable_caching = enable_caching
62 self._verbose = verbose
63 self._cache = dict()
64 self._is_active = True
65 self._all_wrappers = []
66
67 def is_active(self):
68 return self._is_active
69
70 @contextlib.contextmanager
71 def _disable_casts(self):
72 self._is_active = False
73 yield
74 self._is_active = True
75
76 def _clear_cache(self):
77 self._cache.clear()
78
79 # Experimental support for saving / restoring uncasted versions of functions
80 def _save_func(self, mod, fn, func):
81 self._all_wrappers.append((mod, fn, func))
82
83 def _deactivate(self):
84 for mod, fn, func in self._all_wrappers:
85 utils.set_func(mod, fn, func)
86 self._all_wrappers = []
87
88 @property
89 def has_cache(self):
90 return self._enable_caching
91
92 @property
93 def cache(self):
94 return self._cache
95
96 def remove_cache(self, param):
97 if self.has_cache and param in self.cache:
98 del self.cache[param]
99
100 @property
101 def verbose(self):
102 return self._verbose
103
104def _half_helper(verbose=False):
105 def _half_wrapper(self):

Callers 1

initFunction · 0.85

Calls

no outgoing calls

Tested by

no test coverage detected