Decorator for function to support argscope Example: .. code-block:: python from mylib import myfunc myfunc = enable_argscope_for_function(myfunc) Args: func: A function mapping one or multiple tensors to one or multiple tensors.
(func, log_shape=True)
| 67 | |
| 68 | |
| 69 | def enable_argscope_for_function(func, log_shape=True): |
| 70 | """Decorator for function to support argscope |
| 71 | |
| 72 | Example: |
| 73 | |
| 74 | .. code-block:: python |
| 75 | |
| 76 | from mylib import myfunc |
| 77 | myfunc = enable_argscope_for_function(myfunc) |
| 78 | |
| 79 | Args: |
| 80 | func: A function mapping one or multiple tensors to one or multiple |
| 81 | tensors. |
| 82 | log_shape (bool): Specify whether the first input resp. output tensor |
| 83 | shape should be printed once. |
| 84 | |
| 85 | Remarks: |
| 86 | If the function ``func`` returns multiple input or output tensors, |
| 87 | only the first input/output tensor shape is displayed during logging. |
| 88 | |
| 89 | Returns: |
| 90 | The decorated function. |
| 91 | |
| 92 | """ |
| 93 | |
| 94 | assert callable(func), "func should be a callable" |
| 95 | |
| 96 | @wraps(func) |
| 97 | def wrapped_func(*args, **kwargs): |
| 98 | actual_args = copy.copy(get_arg_scope()[func.__name__]) |
| 99 | actual_args.update(kwargs) |
| 100 | out_tensor = func(*args, **actual_args) |
| 101 | in_tensor = args[0] |
| 102 | |
| 103 | ctx = get_current_tower_context() |
| 104 | name = func.__name__ if 'name' not in kwargs else kwargs['name'] |
| 105 | if log_shape: |
| 106 | if ('tower' not in ctx.ns_name.lower()) or ctx.is_main_training_tower: |
| 107 | # we assume the first parameter is the most interesting |
| 108 | if isinstance(out_tensor, tuple): |
| 109 | out_tensor_descr = out_tensor[0] |
| 110 | else: |
| 111 | out_tensor_descr = out_tensor |
| 112 | logger.info("{:<12}: {} --> {}".format( |
| 113 | "'" + name + "'", |
| 114 | get_shape_str(in_tensor), |
| 115 | get_shape_str(out_tensor_descr))) |
| 116 | |
| 117 | return out_tensor |
| 118 | wrapped_func.__argscope_enabled__ = True |
| 119 | return wrapped_func |
| 120 | |
| 121 | |
| 122 | def enable_argscope_for_module(module, log_shape=True): |
no outgoing calls
no test coverage detected