lazy_dyndep.py 2.5 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384
  1. ## @package lazy_dyndep
  2. # Module caffe2.python.lazy_dyndep
  3. import os
  4. from caffe2.python import dyndep, lazy
  5. def RegisterOpsLibrary(name):
  6. """Registers a dynamic library that contains custom operators into Caffe2.
  7. Since Caffe2 uses static variable registration, you can optionally load a
  8. separate .so file that contains custom operators and registers that into
  9. the caffe2 core binary. In C++, this is usually done by either declaring
  10. dependency during compilation time, or via dynload. This allows us to do
  11. registration similarly on the Python side.
  12. Unlike dyndep.InitOpsLibrary, this does not actually parse the c++ file
  13. and refresh operators until caffe2 is called in a fashion which requires
  14. operators. In some large codebases this saves a large amount of time
  15. during import.
  16. It is safe to use within a program that also uses dyndep.InitOpsLibrary
  17. Args:
  18. name: a name that ends in .so, such as "my_custom_op.so". Otherwise,
  19. the command will simply be ignored.
  20. Returns:
  21. None
  22. """
  23. if not os.path.exists(name):
  24. # Note(jiayq): if the name does not exist, instead of immediately
  25. # failing we will simply print a warning, deferring failure to the
  26. # time when an actual call is made.
  27. print('Ignoring {} as it is not a valid file.'.format(name))
  28. return
  29. global _LAZY_IMPORTED_DYNDEPS
  30. _LAZY_IMPORTED_DYNDEPS.add(name)
  31. _LAZY_IMPORTED_DYNDEPS = set()
  32. _error_handler = None
  33. def SetErrorHandler(handler):
  34. """Registers an error handler for errors from registering operators
  35. Since the lazy registration may happen at a much later time, having a dedicated
  36. error handler allows for custom error handling logic. It is highly
  37. recomended to set this to prevent errors from bubbling up in weird parts of the
  38. code.
  39. Args:
  40. handler: a function that takes an exception as a single handler.
  41. Returns:
  42. None
  43. """
  44. global _error_handler
  45. _error_handler = handler
  46. def GetImportedOpsLibraries():
  47. _import_lazy()
  48. return dyndep.GetImportedOpsLibraries()
  49. def _import_lazy():
  50. global _LAZY_IMPORTED_DYNDEPS
  51. if not _LAZY_IMPORTED_DYNDEPS:
  52. return
  53. for name in list(_LAZY_IMPORTED_DYNDEPS):
  54. try:
  55. dyndep.InitOpLibrary(name, trigger_lazy=False)
  56. except BaseException as e:
  57. if _error_handler:
  58. _error_handler(e)
  59. finally:
  60. _LAZY_IMPORTED_DYNDEPS.remove(name)
  61. lazy.RegisterLazyImport(_import_lazy)