def test_args_to_input_spec(self):
a_spec = InputSpec([None, 10], name='a')
b_spec = InputSpec([10], name='b')
a_tensor = paddle.static.data(name='a_var', shape=[4, 10])
b_tensor = paddle.static.data(name='b_var', shape=[4, 10])
kwargs = {'c': 1, 'd': 2}
# case 1
foo_spec = FunctionSpec(foo_func, input_spec=[a_spec, b_spec])
input_with_spec, _ = foo_spec.args_to_input_spec(
(a_tensor, b_tensor, 1, 2), {})
self.assertTrue(len(input_with_spec) == 4)
self.assertTrue(input_with_spec[0] == a_spec) # a
self.assertTrue(input_with_spec[1] == b_spec) # b
self.assertTrue(input_with_spec[2] == 1) # c
self.assertTrue(input_with_spec[3] == 2) # d
# case 2
foo_spec = FunctionSpec(foo_func, input_spec=[a_spec])
input_with_spec, _ = foo_spec.args_to_input_spec((a_tensor, b_tensor),
{})
self.assertTrue(len(input_with_spec) == 2)
self.assertTrue(input_with_spec[0] == a_spec) # a
self.assertTupleEqual(input_with_spec[1].shape, (4, 10)) # b.shape
self.assertEqual(input_with_spec[1].name, 'b_var') # b.name
# case 3
# assert kwargs is None if set `input_spec`
foo_spec = FunctionSpec(foo_func, input_spec=[a_spec])
with self.assertRaises(ValueError):
input_with_spec = foo_spec.args_to_input_spec((a_tensor, b_tensor),
{'c': 4})
# case 4
# assert len(args) >= len(self._input_spec)
foo_spec = FunctionSpec(foo_func, input_spec=[a_spec, b_spec])
with self.assertRaises(ValueError):
input_with_spec = foo_spec.args_to_input_spec((a_tensor, ), {})
class StaticLayer(object):
"""
Wrapper class to Manage program conversion of decorated function.
"""
def __init__(self, function, input_spec=None):
"""
Initializes a `StaticLayer`.
Args:
function(callable): A function or method that will be converted into static program.
input_spec(list[InputSpec]): list of InputSpec to specify the `shape/dtype/name` information for each input argument, default None.
"""
# save the instance `self` while decorating a method of class.
if inspect.ismethod(function):
self._dygraph_function = getattr(function, '__func__')
self._class_instance = getattr(function, '__self__')
else:
self._dygraph_function = function
self._class_instance = None
self._input_spec = input_spec
self._function_spec = FunctionSpec(function, input_spec)
self._program_cache = ProgramCache()
self._descriptor_cache = weakref.WeakKeyDictionary()
# Note: Hold a reference to ProgramTranslator for switching `enable_declarative`.
self._program_trans = ProgramTranslator()
def __get__(self, instance, owner):
"""
Overrides this method to parse the class instance and call bound method correctly.
For example:
'''
class Net(Layer):
def __init__(self):
pass
@paddle.jit.to_static
def forward(self, x, y):
return x + y
net = Net()
out = net(x, y)
'''
In above case, `net(x, y)` will call `net.forward(x, y)` firstly that is a bound method
of `Net` instance. After decorated by `@paddle.jit.to_static`, it will firstly to call `__get__`
to parse the class instance correctly instead of the `StaticLayer` instance.
"""
if instance not in self._descriptor_cache:
if instance is None:
return self
# Note(Aurelius84): To construct new instance of StaticLayer when we
# first encouter the bound function of layer and cache it.
new_static_layer = self._clone()
new_static_layer._class_instance = instance
self._descriptor_cache[instance] = new_static_layer
return self._descriptor_cache[instance]
def _clone(self):
return self.__class__(self._dygraph_function, self._input_spec)
def __call__(self, *args, **kwargs):
"""
Supports to call the returned instance with input `args` and `kwargs` directly.
Args:
*args(tuple): tuple of all input arguments from original decorated function.
**kwargs(dict): dict of all input keyward arguments from original decorated function.
Return:
Outputs of decorated function.
"""
# 1. call dygraph function directly if not enable `declarative`
if not self._program_trans.enable_declarative:
logging_utils.warn(
"The decorator '@paddle.jit.to_static' does NOT work when setting ProgramTranslator.enable=False. "
"We will just return dygraph output.")
return self._call_dygraph_function(*args, **kwargs)
if not in_dygraph_mode() and self._program_trans.enable_declarative:
raise RuntimeError(
"Failed to run the callable object {} decorated by '@paddle.jit.to_static', "
"because it does NOT in dynamic mode. Please disable the static mode to enter dynamic mode with the "
"following API: paddle.disable_static().".format(
self.dygraph_function))
# 2. trace ops from dygraph layers and cache the generated program.
args, kwargs = self._function_spec.unified_args_and_kwargs(args, kwargs)
try:
concrete_program, partial_program_layer = self.get_concrete_program(
*args, **kwargs)
# 3. synchronize self.training attribute.
if isinstance(self._class_instance, layers.Layer):
partial_program_layer.training = self._class_instance.training
# 4. return outputs.
return partial_program_layer(args)
except Exception as e:
if not hasattr(e, ERROR_DATA):
# runtime error
attach_error_data(e, in_runtime=True)
error_data = getattr(e, ERROR_DATA, None)
if error_data:
new_exception = error_data.create_exception()
if six.PY3:
# NOTE(liym27):
# 1. Why `raise new_exception from None`?
# In Python 3, by default, an new exception is raised with trace information of the caught exception.
# This only raises new_exception and hides unwanted implementation details from tracebacks of the
# caught exception.
# 2. Use exec to bypass syntax error checking in Python 2.
six.exec_("raise new_exception from None")
else:
raise new_exception
else:
raise
def _call_dygraph_function(self, *args, **kwargs):
"""
Calls dygraph function directly and returns the outputs.
Args:
*args(tuple): tuple of all input arguments from original decorated function.
**kwargs(dict): dict of all input keyward arguments from original decorated function.
Return:
Outputs of dygraph function.
"""
if self._class_instance is not None:
dygraph_function = self._dygraph_function.__get__(
self._class_instance)
else:
dygraph_function = self._dygraph_function
return dygraph_function(*args, **kwargs)
def get_concrete_program(self, *args, **kwargs):
"""
Returns traced concrete program and inner executable partial layer.
Args:
*args(tuple): input arguments values or InputSpec
**kwargs(dict) : input kwargs values.
Returns:
Traced ConcreteProgram and executable translated Layer.
"""
# 1. unify args/kwargs and replace Tensor with InputSpec
if len(args) != len(self._function_spec.args_name):
args, kwargs = self._function_spec.unified_args_and_kwargs(args,
kwargs)
input_with_spec = self._function_spec.args_to_input_spec(args, kwargs)
# 2. generate cache key
cache_key = CacheKey(self._function_spec, input_with_spec,
self._class_instance)
# 3. check whether hit the cache or build a new program for the input arguments
concrete_program, partial_program_layer = self._program_cache[cache_key]
return concrete_program, partial_program_layer
def get_traced_count(self):
"""
Returns the number of traced programs for the decorated function.
"""
return len(self._program_cache)
@property
def code(self):
"""
Returns the source code of transformed static function for debugging.
"""
static_func = convert_to_static(self._dygraph_function)
source_code = func_to_source_code(static_func)
return source_code
@property
def dygraph_function(self):
"""
Returns the original decorated function.
"""
return self._dygraph_function
@property
def concrete_program(self):
"""
Returns recent ConcreteProgram instance of decorated function.
Examples:
.. code-block:: python
import paddle
from paddle.jit import to_static
from paddle.static import InputSpec
paddle.disable_static()
def foo(x, y):
z = x + y
return z
# usage 1:
decorated_foo = to_static(foo, input_spec=[InputSpec([10], name='x'), InputSpec([10], name='y')])
print(decorated_foo.concrete_program)
# usage 2:
decorated_foo = to_static(foo)
out_foo = decorated_foo(paddle.rand([10]), paddle.rand([10]))
print(decorated_foo.concrete_program)
"""
# if specific the `input_spec`, the length of program_cache will always 1,
# else, return the last one.
cached_program_len = len(self._program_cache)
# If specific `input_spec`, apply convertion from dygraph layers into static Program.
if cached_program_len == 0:
input_spec = self._function_spec.input_spec
has_input_spec = (input_spec is not None and len(input_spec) > 0)
if has_input_spec:
concrete_program, _ = self.get_concrete_program(*input_spec)
return concrete_program
else:
raise ValueError(
"No valid transformed program for {}.\n\t Please specific `input_spec` in `@paddle.jit.to_static` or feed input tensor to call the decorated function at once.\n".
format(self._function_spec))
# If more than one programs have been cached, return the recent converted program by default.
elif cached_program_len > 1:
logging.warning(
"Current {} has more than one cached programs: {}, the last traced progam will be return by default.".
format(self._function_spec, cached_program_len))
cache_key, (concrete_program,
partial_layer) = self._program_cache.last()
return concrete_program
@property
def inputs(self):
"""
Returns input tensors of recent converted static program.
"""
concrete_program = self.concrete_program
inputs = [
var for var in flatten(concrete_program.inputs)
if isinstance(var, framework.Variable)
]
return inputs
@property
def outputs(self):
"""
Returns output tensors of recent converted static program.
"""
concrete_program = self.concrete_program
outputs = [
var for var in flatten(concrete_program.outputs)
if isinstance(var, framework.Variable)
]
return outputs
@property
def main_program(self):
"""
Returns recent converted static main program.
"""
concrete_program = self.concrete_program
main_program = concrete_program.main_program
return main_program
@property
def program_cache(self):
return self._program_cache
@property
def function_spec(self):
return self._function_spec
class StaticFunction(object):
"""
Wrapper class to Manage program conversion of decorated function.
"""
def __init__(self, function, input_spec=None, **kwargs):
"""
Initializes a `StaticFunction`.
Args:
function(callable): A function or method that will be converted into static program.
input_spec(list[InputSpec]): list of InputSpec to specify the `shape/dtype/name` information for each input argument, default None.
**kwargs(dict): other arguments like `build_strategy` et.al.
"""
# save the instance `self` while decorating a method of class.
if inspect.ismethod(function):
self._dygraph_function = getattr(function, '__func__')
self._class_instance = getattr(function, '__self__')
else:
self._dygraph_function = function
self._class_instance = None
self._input_spec = input_spec
self._function_spec = FunctionSpec(function, input_spec)
self._program_cache = ProgramCache()
self._descriptor_cache = weakref.WeakKeyDictionary()
# Note: Hold a reference to ProgramTranslator for switching `enable_to_static`.
self._program_trans = ProgramTranslator()
self._kwargs = kwargs
self._training = True
self._cuda_graph_capture_mode = ""
self._cuda_graph_pool_id = 0
def train(self):
if isinstance(self._class_instance,
layers.Layer) and self._class_instance.training == False:
raise RuntimeError(
"Failed to switch train mode. {} is a Layer's method, "
"please use Layer.train() to switch train mode.".format(
self.dygraph_function))
self._training = True
def eval(self):
if isinstance(self._class_instance,
layers.Layer) and self._class_instance.training == True:
raise RuntimeError(
"Failed to switch eval mode. {} is a Layer's method, "
"please use Layer.eval() to switch eval mode.".format(
self.dygraph_function))
self._training = False
def __get__(self, instance, owner):
"""
Overrides this method to parse the class instance and call bound method correctly.
For example:
'''
class Net(Layer):
def __init__(self):
pass
@paddle.jit.to_static
def forward(self, x, y):
return x + y
net = Net()
out = net(x, y)
'''
In above case, `net(x, y)` will call `net.forward(x, y)` firstly that is a bound method
of `Net` instance. After decorated by `@paddle.jit.to_static`, it will firstly to call `__get__`
to parse the class instance correctly instead of the `StaticFunction` instance.
"""
if instance not in self._descriptor_cache:
if instance is None:
return self
# Note(Aurelius84): To construct new instance of StaticFunction when we
# first encouter the bound function of layer and cache it.
new_static_layer = self._clone()
new_static_layer._class_instance = instance
self._descriptor_cache[instance] = new_static_layer
return self._descriptor_cache[instance]
def _clone(self):
return self.__class__(self._dygraph_function, self._input_spec)
def __call__(self, *args, **kwargs):
"""
Supports to call the returned instance with input `args` and `kwargs` directly.
Args:
*args(tuple): tuple of all input arguments from original decorated function.
**kwargs(dict): dict of all input keyward arguments from original decorated function.
Return:
Outputs of decorated function.
"""
# 1. call dygraph function directly if not enable `declarative`
if not self._program_trans.enable_to_static:
# NOTE(liym27):
# Here calls `warnings.warn` but not `logging_utils.warn` because by default warnings.warn(message)
# will show up **only once**. StaticFunction.__call__ will run many times, it is appropriate to
# display this warning message only once.
logging_utils.warn(
"The decorator '@paddle.jit.to_static' does NOT work when setting ProgramTranslator.enable to False. "
"We will just return dygraph output. If you would like to get static graph output, please call API "
"ProgramTranslator.enable(True)")
return self._call_dygraph_function(*args, **kwargs)
if not _non_static_mode():
raise RuntimeError(
"Failed to run the callable object {} decorated by '@paddle.jit.to_static', "
"because it is NOT in dynamic mode. Please disable the static mode to enter dynamic mode with the "
"following API: paddle.disable_static().".format(
self.dygraph_function))
# 2. trace ops from dygraph layers and cache the generated program.
args, kwargs = self._function_spec.unified_args_and_kwargs(
args, kwargs)
try:
concrete_program, partial_program_layer = self.get_concrete_program(
*args, **kwargs, is_train=self._is_train_mode())
# 3. synchronize self.training attribute.
if isinstance(self._class_instance, layers.Layer):
partial_program_layer.training = self._class_instance.training
else:
partial_program_layer.training = self._training
partial_program_layer._cuda_graph_capture_mode = self._cuda_graph_capture_mode
partial_program_layer._cuda_graph_pool_id = self._cuda_graph_pool_id
# 4. return outputs.
try:
return partial_program_layer(args)
except Exception as e:
if not hasattr(e, error.ERROR_DATA):
# runtime error
error.attach_error_data(e, in_runtime=True)
raise
except Exception as e:
error_data = getattr(e, error.ERROR_DATA, None)
if error_data:
error_data.raise_new_exception()
else:
logging_utils.warn(
"Please file an issue at 'https://github.com/PaddlePaddle/Paddle/issues'"
" if you can't handle this {} yourself.".format(type(e)))
raise e
def _is_train_mode(self):
if self._class_instance is not None:
return self._class_instance.training
else:
return self._training
def _call_dygraph_function(self, *args, **kwargs):
"""
Calls dygraph function directly and returns the outputs.
Args:
*args(tuple): tuple of all input arguments from original decorated function.
**kwargs(dict): dict of all input keyward arguments from original decorated function.
Return:
Outputs of dygraph function.
"""
if self._class_instance is not None:
dygraph_function = self._dygraph_function.__get__(
self._class_instance)
else:
dygraph_function = self._dygraph_function
return dygraph_function(*args, **kwargs)
def get_concrete_program(self, *args, **kwargs):
"""
Returns traced concrete program and inner executable partial layer.
Args:
*args(tuple): input arguments values or InputSpec
**kwargs(dict) : input kwargs values.
Returns:
Traced ConcreteProgram and executable translated Layer.
"""
with_hook = kwargs.get("with_hook", False)
is_train = kwargs.get("is_train", True)
if "is_train" in kwargs: kwargs.pop("is_train")
if "with_hook" in kwargs: kwargs.pop("with_hook")
# 1. unify args/kwargs and replace Tensor with InputSpec
if len(args) != len(self._function_spec.args_name):
args, kwargs = self._function_spec.unified_args_and_kwargs(
args, kwargs)
input_args_with_spec, input_kwargs_with_spec = self._function_spec.args_to_input_spec(
args, kwargs)
# 2. generate cache key
cache_key = CacheKey(self._function_spec,
input_args_with_spec,
input_kwargs_with_spec,
self._class_instance,
**self._kwargs,
with_hook=with_hook,
is_train=is_train)
# 3. check whether hit the cache or build a new program for the input arguments
concrete_program, partial_program_layer = self._program_cache[
cache_key]
return concrete_program, partial_program_layer
def get_traced_count(self):
"""
Returns the number of traced programs for the decorated function.
"""
return len(self._program_cache)
@property
def code(self):
"""
Returns the source code of transformed static function for debugging.
"""
static_func = convert_to_static(self._dygraph_function)
source_code = func_to_source_code(static_func)
return source_code
@property
def dygraph_function(self):
"""
Returns the original decorated function.
"""
return self._dygraph_function
@property
def concrete_program(self):
"""
Returns recent ConcreteProgram instance of decorated function.
Examples:
.. code-block:: python
import paddle
from paddle.jit import to_static
from paddle.static import InputSpec
paddle.disable_static()
def foo(x, y):
z = x + y
return z
# usage 1:
decorated_foo = to_static(foo, input_spec=[InputSpec([10], name='x'), InputSpec([10], name='y')])
print(decorated_foo.concrete_program)
# usage 2:
decorated_foo = to_static(foo)
out_foo = decorated_foo(paddle.rand([10]), paddle.rand([10]))
print(decorated_foo.concrete_program)
"""
return self.concrete_program_specify_input_spec(input_spec=None)
def concrete_program_specify_input_spec(self,
input_spec=None,
with_hook=False):
"""
Returns recent ConcreteProgram instance of decorated function while
specifying input_spec. If the self._function_spec already has
input_spec, it will check the compatibility of input input_spec and
the self._function_spec.input_spec. If input input_spec=None, then
this method uses self._function_spec.input_spec
args:
input_spec (list[InputSpec], optional): Describes the input of
the translate function.
"""
# if specific the `input_spec`, the length of program_cache will always 1,
# else, return the last one.
cached_program_len = len(self._program_cache)
# If specific `input_spec`, apply convertion from dygraph layers into static Program.
if cached_program_len == 0:
desired_input_spec = input_spec
if self._function_spec.input_spec is not None:
if input_spec is not None and not input_specs_compatible(
flatten(input_spec),
flatten(self._function_spec.input_spec)):
raise ValueError(
"The `input_spec`: {} used to construct concrete_program is conflict with the `input_spec`: {} in `@paddle.jit.to_static`"
.format(input_spec, self._function_spec.input_spec))
# NOTE(chenweihang): we should always translated program based on the `input_spec`
# decorated on forward if it is valid
desired_input_spec = self._function_spec.input_spec
if input_spec is not None:
logging_utils.warn(
"\n\nYou have specified `input_spec` both in function definition (higher priority) and `paddle.jit.save` (will be ignored.)\n\n\t Using: {}\n\n\t Ignore: {}\n"
.format(desired_input_spec, input_spec))
has_input_spec = (desired_input_spec is not None)
if has_input_spec:
concrete_program, _ = self.get_concrete_program(
*desired_input_spec,
with_hook=with_hook,
is_train=self._is_train_mode())
return concrete_program
else:
raise ValueError(
"No valid transformed program for {}.\n\t Please specific `input_spec` in `@paddle.jit.to_static` or feed input tensor to call the decorated function at once.\n"
.format(self._function_spec))
elif with_hook:
cache_key = self._program_cache._recent_cache_key
cache_key.kwargs["with_hook"] = True
concrete_program, _ = self._program_cache[cache_key]
return concrete_program
# If more than one programs have been cached, return the recent converted program by default.
elif cached_program_len > 1:
logging_utils.warn(
"Current {} has more than one cached programs: {}, the last traced progam will be return by default."
.format(self._function_spec, cached_program_len))
cache_key, (concrete_program,
partial_layer) = self._program_cache.last()
return concrete_program
@property
def inputs(self):
"""
Returns input tensors of recent converted static program.
"""
concrete_program = self.concrete_program
inputs = [
var for var in flatten(concrete_program.inputs)
if isinstance(var, framework.Variable)
]
return inputs
@property
def outputs(self):
"""
Returns output tensors of recent converted static program.
"""
concrete_program = self.concrete_program
outputs = [
var for var in flatten(concrete_program.outputs)
if isinstance(var, framework.Variable)
]
return outputs
@property
def main_program(self):
"""
Returns recent converted static main program.
"""
concrete_program = self.concrete_program
main_program = concrete_program.main_program
return main_program
@property
def program_cache(self):
return self._program_cache
@property
def function_spec(self):
return self._function_spec