mirror of
https://github.com/zebrajr/pytorch.git
synced 2025-12-07 12:21:27 +01:00
07fd5f8ff9
Summary: Pull Request resolved: https://github.com/pytorch/pytorch/pull/39488 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). 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. Differential Revision: D21870844 fbshipit-source-id: 3f65fedb65bb48663670349cee5e1d3e22d560ed
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)
|