def _wrap_op(op_class, submodule, parent_module, wrapper_doc):
"""Wrap the DALI Operator with fn API and insert the function into appropriate module.
Args:
op_class: Op class to wrap
submodule: Additional submodule (scope)
parent_module (str): If set to None, the wrapper is placed in nvidia.dali.fn module,
otherwise in a specified parent module.
wrapper_doc (str): Documentation of the wrapper function
"""
schema = _b.TryGetSchema(op_class.__name__)
make_hidden = schema.IsDocHidden() if schema else False
wrapper_name = _to_snake_case(op_class.__name__)
if parent_module is None:
fn_module = sys.modules[__name__]
else:
fn_module = sys.modules[parent_module]
module = _internal.get_submodule(fn_module, submodule)
if not hasattr(module, wrapper_name):
wrap_func = _wrap_op_fn(op_class, wrapper_name, wrapper_doc)
setattr(module, wrapper_name, wrap_func)
if submodule:
wrap_func.__module__ = module.__name__
if make_hidden:
parent_module = _internal.get_submodule(fn_module, submodule[:-1])
setattr(parent_module, wrapper_name, wrap_func)
def _load_ops():
global _cpu_ops
global _gpu_ops
global _mixed_ops
_cpu_ops = _cpu_ops.union(set(_b.RegisteredCPUOps()))
_gpu_ops = _gpu_ops.union(set(_b.RegisteredGPUOps()))
_mixed_ops = _mixed_ops.union(set(_b.RegisteredMixedOps()))
_cpu_gpu_ops = _cpu_ops.union(_gpu_ops).union(_mixed_ops)
ops_module = sys.modules[__name__]
for op_reg_name in _cpu_gpu_ops:
schema = _b.TryGetSchema(op_reg_name)
make_hidden = schema.IsDocHidden() if schema else False
op_full_name, submodule, op_name = _process_op_name(op_reg_name, make_hidden)
module = _internal.get_submodule(ops_module, submodule)
if not hasattr(module, op_name):
op_class = python_op_factory(op_name, op_reg_name, op_device = "cpu")
op_class.__module__ = module.__name__
setattr(module, op_name, op_class)
if op_name not in ["ExternalSource"]:
_wrap_op(op_class, submodule)
# The operator was inserted into nvidia.dali.ops.hidden module, let's import it here
# so it would be usable, but not documented as coming from other module
if make_hidden:
parent_module = _internal.get_submodule(ops_module, submodule[:-1])
setattr(parent_module, op_name, op_class)
def _wrap_eager_op(op_class, submodule, wrapper_name, wrapper_doc):
"""Exposes eager operator to the appropriate module (similar to :func:`nvidia.dali.fn._wrap_op`).
Uses ``op_class`` for preprocessing inputs and keyword arguments and filling OpSpec for backend
eager operators.
Args:
op_class: Op class to wrap.
submodule: Additional submodule (scope).
wrapper_name: Wrapper name (the same as in fn API).
wrapper_doc (str): Documentation of the wrapper function.
"""
op_name = op_class.schema_name
op_schema = _b.TryGetSchema(op_name)
if op_schema.IsDeprecated(
) or op_name in _stateful_operators or op_name in _generator_operators:
# TODO(ksztenderski): For now only exposing stateless operators.
return
else:
# If operator is not stateful or a generator expose it as stateless.
wrapper = _wrap_stateless(op_class, op_name, wrapper_name)
# Exposing to eager.experimental module.
eager_module = _internal.get_submodule(sys.modules[__name__],
'experimental')
op_module = _internal.get_submodule(eager_module, submodule)
if not hasattr(op_module, wrapper_name):
wrapper.__name__ = wrapper_name
wrapper.__qualname__ = wrapper_name
wrapper.__doc__ = wrapper_doc
if submodule:
wrapper.__module__ = op_module.__name__
setattr(op_module, wrapper_name, wrapper)
def _wrap_op(op_class, submodule):
schema = _b.TryGetSchema(op_class.__name__)
make_hidden = schema.IsDocHidden() if schema else False
wrapper_name = _to_snake_case(op_class.__name__)
fn_module = sys.modules[__name__]
module = _internal.get_submodule(fn_module, submodule)
if not hasattr(module, wrapper_name):
wrap_func = _wrap_op_fn(op_class, wrapper_name)
setattr(module, wrapper_name, wrap_func)
if submodule:
wrap_func.__module__ = module.__name__
if make_hidden:
parent_module = _internal.get_submodule(fn_module, submodule[:-1])
setattr(parent_module, wrapper_name, wrap_func)
def _load_readers_tfrecord():
_TFRecordReaderImpl.__call__.__doc__ = _docstring_generator_call(
"readers__TFRecord")
global _cpu_ops
_cpu_ops = _cpu_ops.union({'readers__TFRecord', 'TFRecordReader'})
ops_module = sys.modules[__name__]
class TFRecordReader(_TFRecordReaderImpl, metaclass=_DaliOperatorMeta):
pass
class TFRecord(_TFRecordReaderImpl, metaclass=_DaliOperatorMeta):
pass
for op_reg_name, internal_schema, op_class in [
('readers__TFRecord', 'readers___TFRecord', TFRecord),
('TFRecordReader', '_TFRecordReader', TFRecordReader)
]:
op_class.schema_name = op_reg_name
op_class._internal_schema_name = internal_schema
op_full_name, submodule, op_name = _process_op_name(op_reg_name)
module = _internal.get_submodule(ops_module, submodule)
if not hasattr(module, op_name):
op_class.__module__ = module.__name__
setattr(module, op_name, op_class)
_wrap_op(op_class, submodule)
def _get_eager_target_module(parent_module, submodules, make_hidden):
""" Returns target module inside ``parent_module`` if specified, otherwise inside eager. """
if parent_module is None:
# Exposing to nvidia.dali.experimental.eager module.
parent_module = _internal.get_submodule('nvidia.dali',
'experimental.eager')
else:
# Exposing to experimental.eager submodule of the specified parent module.
parent_module = _internal.get_submodule(sys.modules[parent_module],
'experimental.eager')
if make_hidden:
op_module = _internal.get_submodule(parent_module, submodules[:-1])
else:
op_module = _internal.get_submodule(parent_module, submodules)
return op_module
def _wrap_op(op_class, submodule):
wrapper_name = _to_snake_case(op_class.__name__)
fn_module = sys.modules[__name__]
module = _internal.get_submodule(fn_module, submodule)
if not hasattr(module, wrapper_name):
wrap_func = _wrap_op_fn(op_class, wrapper_name)
setattr(module, wrapper_name, wrap_func)
if submodule:
wrap_func.__module__ = module.__name__
def _expose_eager_op_as_object(op_class, submodule):
""" Exposes eager operators as objects. Can be used if we decide to change eager API from
functional to objective.
"""
op_name = op_class.schema_name
module = _internal.get_submodule('nvidia.dali.experimental.eager',
submodule)
op = _eager_op_object_factory(op_class, op_name)
setattr(module, op_name, op)
def _load_ops():
global _cpu_ops
global _gpu_ops
global _mixed_ops
_cpu_ops = _cpu_ops.union(set(_b.RegisteredCPUOps()))
_gpu_ops = _gpu_ops.union(set(_b.RegisteredGPUOps()))
_mixed_ops = _mixed_ops.union(set(_b.RegisteredMixedOps()))
_cpu_gpu_ops = _cpu_ops.union(_gpu_ops).union(_mixed_ops)
ops_module = sys.modules[__name__]
for op_reg_name in _cpu_gpu_ops:
op_full_name, submodule, op_name = _process_op_name(op_reg_name)
module = _internal.get_submodule(ops_module, submodule)
if not hasattr(module, op_name):
op_class = python_op_factory(op_name, op_reg_name, op_device="cpu")
op_class.__module__ = module.__name__
setattr(module, op_name, op_class)
if op_name not in ["ExternalSource"]:
_wrap_op(op_class, submodule)
def _get_experimental():
# TODO(klecki): this is WAR only for experimental module
current_module = sys.modules[__name__]
experimental = _internal.get_submodule(current_module, "experimental")
return experimental
