dyndep.py 1.5 KB

12345678910111213141516171819202122232425262728293031323334353637383940414243444546474849505152
  1. ## @package dyndep
  2. # Module caffe2.python.dyndep
  3. import ctypes
  4. import os
  5. from threading import Lock
  6. from caffe2.python import core, extension_loader
  7. def InitOpsLibrary(name, trigger_lazy=True):
  8. """Loads a dynamic library that contains custom operators into Caffe2.
  9. Since Caffe2 uses static variable registration, you can optionally load a
  10. separate .so file that contains custom operators and registers that into
  11. the caffe2 core binary. In C++, this is usually done by either declaring
  12. dependency during compilation time, or via dynload. This allows us to do
  13. registration similarly on the Python side.
  14. Args:
  15. name: a name that ends in .so, such as "my_custom_op.so". Otherwise,
  16. the command will simply be ignored.
  17. Returns:
  18. None
  19. """
  20. if not os.path.exists(name):
  21. # Note(jiayq): if the name does not exist, instead of immediately
  22. # failing we will simply print a warning, deferring failure to the
  23. # time when an actual call is made.
  24. print('Ignoring {} as it is not a valid file.'.format(name))
  25. return
  26. _init_impl(name, trigger_lazy=trigger_lazy)
  27. _IMPORTED_DYNDEPS = set()
  28. dll_lock = Lock()
  29. def GetImportedOpsLibraries():
  30. return _IMPORTED_DYNDEPS
  31. def _init_impl(path, trigger_lazy=True):
  32. with dll_lock:
  33. _IMPORTED_DYNDEPS.add(path)
  34. with extension_loader.DlopenGuard():
  35. ctypes.CDLL(path)
  36. # reinitialize available ops
  37. core.RefreshRegisteredOperators(trigger_lazy)