mirror of
https://github.com/zebrajr/pytorch.git
synced 2025-12-06 12:20:52 +01:00
415ff0bceb
Summary: Pull Request resolved: https://github.com/pytorch/pytorch/pull/41343 Currently caffe2.InitOpLibrary does the dll import uniliaterally. Instead if we make a lazy version and use it, then many pieces of code which do not need the caffe2urrenoperators get a lot faster. One a real test, the import time went from 140s to 68s. 8s. This also cleans up the algorithm slightly (although it makes a very minimal difference), by parsing the list of operators once, rather than every time a new operator is added, since we defer the RefreshCall until after we've imported all the operators. The key way we maintain safety, is that as soon as someone does an operation which requires a operator (or could), we force importing of all available operators. Future work could include trying to identify which code is needed for which operator and only import the needed ones. There may also be wins available by playing with dlmopen (which opens within a namespace), or seeing if the dl flags have an impact (I tried this and didn't see an impact, but dlmopen may make it better). Note that this was previously landed and reverted. The issue was that if a import failed and raised an exception, the specific library would not be removed from the lazy imports. This caused our tests which had libraries that failed to poison all other tests that ran after it. This has been fixed and a unit test has been added for this case (to help make it obvious what failed). Test Plan: I added a new test a lazy_dyndep_test.py (copied from all_compare_test.py). I'm a little concerned that I don't see any explicit tests for dyndep, but this should provide decent coverage. I've added a specific test to handle the poisoning issues mentioned above, which caused the previous version to get reverted. Differential Revision: D22506369 fbshipit-source-id: 7395df4778e8eb0220630c570360b99a7d60eb83
85 lines
2.6 KiB
Python
85 lines
2.6 KiB
Python
## @package lazy_dyndep
|
|
# Module caffe2.python.lazy_dyndep
|
|
from __future__ import absolute_import
|
|
from __future__ import division
|
|
from __future__ import print_function
|
|
from __future__ import unicode_literals
|
|
|
|
import os
|
|
from caffe2.python import core, dyndep
|
|
|
|
|
|
def RegisterOpsLibrary(name):
|
|
"""Registers a dynamic library that contains custom operators into Caffe2.
|
|
|
|
Since Caffe2 uses static variable registration, you can optionally load a
|
|
separate .so file that contains custom operators and registers that into
|
|
the caffe2 core binary. In C++, this is usually done by either declaring
|
|
dependency during compilation time, or via dynload. This allows us to do
|
|
registration similarly on the Python side.
|
|
|
|
Unlike dyndep.InitOpsLibrary, this does not actually parse the c++ file
|
|
and refresh operators until caffe2 is called in a fashion which requires
|
|
operators. In some large codebases this saves a large amount of time
|
|
during import.
|
|
|
|
It is safe to use within a program that also uses dyndep.InitOpsLibrary
|
|
|
|
Args:
|
|
name: a name that ends in .so, such as "my_custom_op.so". Otherwise,
|
|
the command will simply be ignored.
|
|
Returns:
|
|
None
|
|
"""
|
|
if not os.path.exists(name):
|
|
# Note(jiayq): if the name does not exist, instead of immediately
|
|
# failing we will simply print a warning, deferring failure to the
|
|
# time when an actual call is made.
|
|
print('Ignoring {} as it is not a valid file.'.format(name))
|
|
return
|
|
global _LAZY_IMPORTED_DYNDEPS
|
|
_LAZY_IMPORTED_DYNDEPS.add(name)
|
|
|
|
|
|
_LAZY_IMPORTED_DYNDEPS = set()
|
|
_error_handler = None
|
|
|
|
|
|
def SetErrorHandler(handler):
|
|
"""Registers an error handler for errors from registering operators
|
|
|
|
Since the lazy registration may happen at a much later time, having a dedicated
|
|
error handler allows for custom error handling logic. It is highly
|
|
recomended to set this to prevent errors from bubbling up in weird parts of the
|
|
code.
|
|
|
|
Args:
|
|
handler: a function that takes an exception as a single handler.
|
|
Returns:
|
|
None
|
|
"""
|
|
|
|
global _error_handler
|
|
_error_handler = handler
|
|
|
|
|
|
def GetImportedOpsLibraries():
|
|
_import_lazy()
|
|
return dyndep.GetImportedOpsLibraries()
|
|
|
|
|
|
def _import_lazy():
|
|
global _LAZY_IMPORTED_DYNDEPS
|
|
if not _LAZY_IMPORTED_DYNDEPS:
|
|
return
|
|
for name in list(_LAZY_IMPORTED_DYNDEPS):
|
|
try:
|
|
dyndep.InitOpLibrary(name, trigger_lazy=False)
|
|
except BaseException as e:
|
|
if _error_handler:
|
|
_error_handler(e)
|
|
finally:
|
|
_LAZY_IMPORTED_DYNDEPS.remove(name)
|
|
|
|
core.RegisterLazyImport(_import_lazy)
|