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
53 lines
1.6 KiB
Python
53 lines
1.6 KiB
Python
## @package dyndep
|
|
# Module caffe2.python.dyndep
|
|
from __future__ import absolute_import
|
|
from __future__ import division
|
|
from __future__ import print_function
|
|
from __future__ import unicode_literals
|
|
|
|
import ctypes
|
|
import os
|
|
from threading import Lock
|
|
from caffe2.python import core, extension_loader
|
|
|
|
|
|
def InitOpsLibrary(name, trigger_lazy=True):
|
|
"""Loads 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.
|
|
|
|
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
|
|
_init_impl(name, trigger_lazy=trigger_lazy)
|
|
|
|
|
|
_IMPORTED_DYNDEPS = set()
|
|
dll_lock = Lock()
|
|
|
|
|
|
def GetImportedOpsLibraries():
|
|
return _IMPORTED_DYNDEPS
|
|
|
|
|
|
def _init_impl(path, trigger_lazy=True):
|
|
with dll_lock:
|
|
_IMPORTED_DYNDEPS.add(path)
|
|
with extension_loader.DlopenGuard():
|
|
ctypes.CDLL(path)
|
|
# reinitialize available ops
|
|
core.RefreshRegisteredOperators(trigger_lazy)
|