| 123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384 |
- ## @package lazy_dyndep
- # Module caffe2.python.lazy_dyndep
- import os
- from caffe2.python import dyndep, lazy
- 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)
- lazy.RegisterLazyImport(_import_lazy)
|