mirror of
https://github.com/zebrajr/pytorch.git
synced 2025-12-06 12:20:52 +01:00
dfa914a90c
Summary: Pull Request resolved: https://github.com/pytorch/pytorch/pull/41687 Specifically, this makes a new library (lazy), which can be used from both core and workspace. This allows workspace.Createnet to trigger lazy loading of dyndep dependencies. Test Plan: Added a unit test specifically for workspace.CreateNet Reviewed By: dzhulgakov Differential Revision: D22441877 fbshipit-source-id: 3a9d1af9962585d08ea2566c9c85bec7377d39f2
85 lines
2.6 KiB
Python
85 lines
2.6 KiB
Python
## @package lazy_dyndep
|
|
# Module caffe2.python.lazy_dyndep
|
|
from __future__ import absolute_import
|
|
from __future__ import division
|
|
from __future__ import print_function
|
|
from __future__ import unicode_literals
|
|
|
|
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)
|