wehub-resource-sync
1762 lines
67 KiB
Python
1762 lines
67 KiB
Python
# Copyright (c) 2019 PaddlePaddle Authors. All Rights Reserved.
|
|
# Copyright (c) 2021 NVIDIA Corporation. All rights reserved.
|
|
#
|
|
# Licensed under the Apache License, Version 2.0 (the "License");
|
|
# you may not use this file except in compliance with the License.
|
|
# You may obtain a copy of the License at
|
|
#
|
|
# http://www.apache.org/licenses/LICENSE-2.0
|
|
#
|
|
# Unless required by applicable law or agreed to in writing, software
|
|
# distributed under the License is distributed on an "AS IS" BASIS,
|
|
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
|
# See the License for the specific language governing permissions and
|
|
# limitations under the License.
|
|
|
|
from __future__ import annotations
|
|
|
|
import inspect
|
|
import os
|
|
import pickle
|
|
import sys
|
|
import threading
|
|
import types
|
|
import warnings
|
|
from collections import OrderedDict
|
|
from collections.abc import Callable, Sequence
|
|
from contextlib import contextmanager
|
|
from types import ModuleType
|
|
from typing import (
|
|
TYPE_CHECKING,
|
|
Any,
|
|
Literal,
|
|
Protocol,
|
|
TypeAlias,
|
|
TypedDict,
|
|
TypeVar,
|
|
overload,
|
|
)
|
|
|
|
from typing_extensions import (
|
|
NotRequired,
|
|
ParamSpec,
|
|
Unpack,
|
|
)
|
|
|
|
import paddle
|
|
from paddle.base import core, dygraph
|
|
from paddle.base.compiler import (
|
|
BuildStrategy,
|
|
)
|
|
from paddle.base.dygraph.base import (
|
|
switch_to_static_graph,
|
|
)
|
|
from paddle.base.executor import Executor, scope_guard
|
|
from paddle.base.framework import (
|
|
EagerParamBase,
|
|
Parameter,
|
|
Variable,
|
|
_current_expected_place,
|
|
dygraph_only,
|
|
)
|
|
from paddle.base.wrapped_decorator import wrap_decorator
|
|
from paddle.framework import use_pir_api
|
|
from paddle.nn import Layer
|
|
from paddle.static.io import save_inference_model
|
|
from paddle.utils.environments import (
|
|
EnvironmentVariableGuard,
|
|
)
|
|
|
|
from .dy2static import logging_utils
|
|
from .dy2static.convert_call_func import add_ignore_module
|
|
from .dy2static.program_translator import (
|
|
ASTStaticFunction,
|
|
ProgramTranslator,
|
|
StaticFunction,
|
|
SymbolicStaticFunction,
|
|
unwrap_decorators,
|
|
)
|
|
from .dy2static.utils import (
|
|
ENV_ENABLE_SOT,
|
|
Backend,
|
|
infer_use_cinn_backend,
|
|
)
|
|
from .pir_translated_layer import PIR_INFER_MODEL_SUFFIX, PirTranslatedLayer
|
|
from .translated_layer import (
|
|
INFER_MODEL_SUFFIX,
|
|
INFER_PARAMS_INFO_SUFFIX,
|
|
INFER_PARAMS_SUFFIX,
|
|
INFER_PROPERTY_SUFFIX,
|
|
TranslatedLayer,
|
|
)
|
|
|
|
if TYPE_CHECKING:
|
|
from paddle import Tensor
|
|
from paddle._typing import NestedStructure
|
|
from paddle.static import InputSpec
|
|
|
|
class _SaveOptions(TypedDict):
|
|
output_spec: NotRequired[Sequence[Tensor | int]]
|
|
with_hook: NotRequired[bool]
|
|
combine_params: NotRequired[bool]
|
|
clip_extra: NotRequired[bool]
|
|
skip_forward: NotRequired[bool]
|
|
input_names_after_prune: NotRequired[list[str]]
|
|
skip_prune_program: NotRequired[bool]
|
|
separate_parameters: NotRequired[bool]
|
|
|
|
class _LoadOptions(TypedDict):
|
|
model_filename: NotRequired[str]
|
|
params_filename: NotRequired[str]
|
|
|
|
|
|
_LayerT = TypeVar("_LayerT", bound=Layer)
|
|
_RetT = TypeVar("_RetT")
|
|
_InputT = ParamSpec("_InputT")
|
|
Backends: TypeAlias = Literal["CINN"]
|
|
|
|
|
|
@contextmanager
|
|
def sot_mode_guard(value: bool):
|
|
with EnvironmentVariableGuard(ENV_ENABLE_SOT, value):
|
|
yield
|
|
|
|
|
|
def copy_decorator_attrs(original_func, decorated_obj):
|
|
"""
|
|
Copies some necessary attributes from original function into decorated function.
|
|
|
|
Args:
|
|
original_func(callable): the original decorated function.
|
|
decorated_obj(StaticFunction): the target decorated StaticFunction object.
|
|
"""
|
|
decorator_name = "to_static"
|
|
|
|
decorated_obj.__name__ = original_func.__name__
|
|
decorated_obj._decorator_name = decorator_name
|
|
if not inspect.ismethod(original_func):
|
|
decorated_obj.__wrapped__ = original_func
|
|
decorated_obj.__doc__ = original_func.__doc__
|
|
if hasattr(original_func, "__module__"):
|
|
decorated_obj.__module__ = original_func.__module__
|
|
|
|
return decorated_obj
|
|
|
|
|
|
def ignore_module(modules: list[ModuleType]) -> None:
|
|
"""
|
|
Adds modules that ignore transcription.
|
|
Builtin modules that have been ignored are collections, pdb, copy, inspect, re, numpy, logging, six
|
|
|
|
Args:
|
|
modules (list[ModuleType]): Ignored modules that you want to add
|
|
|
|
Examples:
|
|
.. code-block:: pycon
|
|
|
|
>>> import scipy
|
|
>>> import networkx
|
|
|
|
>>> import paddle
|
|
>>> from paddle.jit import ignore_module
|
|
>>> modules = [
|
|
... scipy,
|
|
... networkx,
|
|
... ]
|
|
>>> ignore_module(modules)
|
|
|
|
"""
|
|
add_ignore_module(modules)
|
|
|
|
|
|
class _ToStaticOptions(TypedDict):
|
|
property: NotRequired[bool]
|
|
full_graph: NotRequired[bool]
|
|
|
|
|
|
class _ToStaticDecorator(Protocol):
|
|
@overload
|
|
def __call__(self, function: _LayerT) -> _LayerT: ...
|
|
|
|
@overload
|
|
def __call__(
|
|
self, function: Callable[_InputT, _RetT]
|
|
) -> StaticFunction[_InputT, _RetT]: ...
|
|
|
|
|
|
@overload
|
|
def to_static(
|
|
function: _LayerT,
|
|
input_spec: NestedStructure[InputSpec] | None = ...,
|
|
build_strategy: BuildStrategy | None = ...,
|
|
backend: Backends | None = ...,
|
|
**kwargs: Unpack[_ToStaticOptions],
|
|
) -> _LayerT: ...
|
|
|
|
|
|
@overload
|
|
def to_static(
|
|
function: Callable[_InputT, _RetT],
|
|
input_spec: NestedStructure[InputSpec] | None = ...,
|
|
build_strategy: BuildStrategy | None = ...,
|
|
backend: Backends | None = ...,
|
|
**kwargs: Unpack[_ToStaticOptions],
|
|
) -> StaticFunction[_InputT, _RetT]: ...
|
|
|
|
|
|
@overload
|
|
def to_static(
|
|
function: None = ...,
|
|
input_spec: NestedStructure[InputSpec] | None = ...,
|
|
build_strategy: BuildStrategy | None = ...,
|
|
backend: Backends | None = ...,
|
|
**kwargs: Unpack[_ToStaticOptions],
|
|
) -> _ToStaticDecorator: ...
|
|
|
|
|
|
def to_static(
|
|
function=None,
|
|
input_spec=None,
|
|
build_strategy=None,
|
|
backend="CINN",
|
|
**kwargs,
|
|
):
|
|
"""
|
|
Converts dynamic graph APIs into static graph function APIs. Decorator
|
|
@to_static handles the Program and Executor of static graph mode and returns
|
|
the result as dynamic graph Tensor(s). Users could use the returned dynamic
|
|
graph Tensor(s) to do dynamic graph training, inference, or other operations.
|
|
If the decorated function calls other dynamic graph function, the called one
|
|
will be converted into static graph function as well.
|
|
|
|
Args:
|
|
function (callable): Callable dynamic graph function. If it used as a
|
|
decorator, the decorated function will be parsed as this parameter.
|
|
input_spec (list[InputSpec]|tuple[InputSpec]): list/tuple of InputSpec to
|
|
specific the shape/dtype/name information of each input Tensor.
|
|
build_strategy (BuildStrategy|None): This argument is used to compile the
|
|
converted program with the specified options, such as operators' fusion
|
|
in the computational graph and memory optimization during the execution
|
|
of the computational graph. For more information about :attr:`build_strategy`,
|
|
please refer to :ref:`paddle.static.BuildStrategy <cn_api_paddle_static_BuildStrategy>`.
|
|
The default is ``None``.
|
|
backend(str, Optional): Specifies compilation backend, which can be ``"CINN"`` or
|
|
``None``. When backend is ``"CINN"``, CINN compiler will be used to speed up
|
|
training and inference. default value is ``"CINN"``.
|
|
kwargs: Support keys including :attr:`property` and :attr:`full_graph`.
|
|
|
|
- property (bool): If True, the function will be treated as a property
|
|
function. The default is False.
|
|
- full_graph (bool): If True, the function will be converted into a
|
|
full static graph. The default is False.
|
|
|
|
Returns:
|
|
Tensor(s): containing the numerical result.
|
|
|
|
Examples:
|
|
.. code-block:: pycon
|
|
|
|
>>> # doctest: +SKIP('`paddle.jit.to_static` can not run in xdoctest')
|
|
>>> import paddle
|
|
>>> from paddle.jit import to_static
|
|
|
|
>>> @to_static
|
|
>>> def func(x):
|
|
... if paddle.mean(x) < 0:
|
|
... x_v = x - 1
|
|
... else:
|
|
... x_v = x + 1
|
|
... return x_v
|
|
>>> x = paddle.ones([1, 2], dtype='float32')
|
|
>>> x_v = func(x)
|
|
>>> print(x_v)
|
|
Tensor(shape=[1, 2], dtype=float32, place=Place(cpu), stop_gradient=True,
|
|
[[2., 2.]])
|
|
|
|
"""
|
|
property = kwargs.get("property", False)
|
|
full_graph = kwargs.get("full_graph", None)
|
|
build_strategy = build_strategy or BuildStrategy()
|
|
if not isinstance(build_strategy, BuildStrategy):
|
|
raise TypeError(
|
|
f"Required type(build_strategy) shall be `paddle.static.BuildStrategy`, but received {type(build_strategy).__name__}"
|
|
)
|
|
backend = Backend.from_arg(backend)
|
|
if infer_use_cinn_backend(backend, build_strategy):
|
|
backend = Backend.CINN
|
|
elif backend.is_pcc():
|
|
pass
|
|
else:
|
|
backend = Backend.PHI
|
|
|
|
def decorated(python_func):
|
|
"""
|
|
Decorates a python function into a ASTStaticFunction or SymbolicStaticFunction object.
|
|
"""
|
|
|
|
nonlocal full_graph
|
|
if full_graph is None:
|
|
flag = ENV_ENABLE_SOT.get()
|
|
full_graph = not flag
|
|
|
|
if sys.version_info >= (3, 15) and not full_graph:
|
|
warnings.warn(
|
|
"full_graph=False is not supported in Python 3.15+. Set full_graph=True automatically"
|
|
)
|
|
full_graph = True
|
|
|
|
StaticClass = (
|
|
ASTStaticFunction if full_graph else SymbolicStaticFunction
|
|
)
|
|
|
|
# Step 1. unwrap the function if it is already decorated.
|
|
_, python_func = unwrap_decorators(python_func)
|
|
|
|
# Step 2. copy some attributes from original python function.
|
|
static_layer = copy_decorator_attrs(
|
|
original_func=python_func,
|
|
decorated_obj=StaticClass(
|
|
function=python_func,
|
|
input_spec=input_spec,
|
|
build_strategy=build_strategy,
|
|
property=property,
|
|
backend=backend,
|
|
),
|
|
)
|
|
|
|
return static_layer
|
|
|
|
# for usage: `to_static(foo, ...)`
|
|
if function is not None:
|
|
if isinstance(function, Layer):
|
|
if isinstance(function.forward, StaticFunction):
|
|
class_name = function.__class__.__name__
|
|
logging_utils.warn(
|
|
f"`{class_name}.forward` has already been decorated somewhere. It will be redecorated to replace previous one."
|
|
)
|
|
function._original_funcs["forward"] = function.forward
|
|
function.forward = decorated(function.forward)
|
|
return function
|
|
else:
|
|
return decorated(function)
|
|
|
|
# for usage: `@to_static`
|
|
return decorated
|
|
|
|
|
|
class _SaveLoadConfig:
|
|
def __init__(self):
|
|
self._output_spec = None
|
|
self._model_filename = None
|
|
self._params_filename = None
|
|
self._separate_params = False
|
|
# used for `paddle.load`
|
|
self._keep_name_table = False
|
|
|
|
# NOTE: Users rarely use following configs, so these configs are not open to users,
|
|
# reducing user learning costs, but we retain the configuration capabilities
|
|
|
|
# If True, programs are modified to only support direct inference deployment.
|
|
# Otherwise,more information will be stored for flexible optimization and re-training.
|
|
# Currently, only True is supported
|
|
self._export_for_deployment = True
|
|
|
|
# If True, It will save inference program only, and do not save params of Program
|
|
self._program_only = False
|
|
self.with_hook = False
|
|
|
|
# if True, multi `StaticFunction` will share params in one file.
|
|
self.combine_params = False
|
|
|
|
# when need to save a prune model, use input_names_after_prune to specify the inputs left after pruning
|
|
self.input_names_after_prune = None
|
|
|
|
# in the scene of llm-inference, pruning program can cause unexpectable result, an option to skip prune is necessary
|
|
self.skip_prune_program = False
|
|
|
|
# if True, the params will be saved separately in multiple files.
|
|
self.separate_parameters = False
|
|
|
|
@property
|
|
def output_spec(self):
|
|
return self._output_spec
|
|
|
|
@output_spec.setter
|
|
def output_spec(self, spec):
|
|
if spec is None:
|
|
return
|
|
if not isinstance(spec, list):
|
|
raise TypeError(
|
|
f"The config `output_spec` should be 'list', but received input type is {type(input)}."
|
|
)
|
|
for var in spec:
|
|
if not isinstance(var, core.eager.Tensor):
|
|
raise TypeError(
|
|
f"The element in config `output_spec` list should be 'Variable', but received element's type is {type(var)}."
|
|
)
|
|
self._output_spec = spec
|
|
|
|
@property
|
|
def model_filename(self):
|
|
return self._model_filename
|
|
|
|
@model_filename.setter
|
|
def model_filename(self, filename):
|
|
if filename is None:
|
|
return
|
|
if not isinstance(filename, str):
|
|
raise TypeError(
|
|
f"The config `model_filename` should be str, but received input's type is {type(filename)}."
|
|
)
|
|
if len(filename) == 0:
|
|
raise ValueError("The config `model_filename` is empty string.")
|
|
self._model_filename = filename
|
|
|
|
@property
|
|
def params_filename(self):
|
|
return self._params_filename
|
|
|
|
@params_filename.setter
|
|
def params_filename(self, filename):
|
|
if filename is None:
|
|
return
|
|
if not isinstance(filename, str):
|
|
raise TypeError(
|
|
f"The config `params_filename` should be str, but received input's type is {type(filename)}."
|
|
)
|
|
if len(filename) == 0:
|
|
raise ValueError("The config `params_filename` is empty string.")
|
|
self._params_filename = filename
|
|
|
|
@property
|
|
def keep_name_table(self):
|
|
return self._keep_name_table
|
|
|
|
@keep_name_table.setter
|
|
def keep_name_table(self, value):
|
|
if value is None:
|
|
return
|
|
if not isinstance(value, bool):
|
|
raise TypeError(
|
|
f"The config `keep_name_table` should be bool value, but received input's type is {type(value)}."
|
|
)
|
|
self._keep_name_table = value
|
|
|
|
|
|
def _parse_save_configs(configs: _SaveOptions) -> _SaveLoadConfig:
|
|
supported_configs = [
|
|
"output_spec",
|
|
"with_hook",
|
|
"combine_params",
|
|
"clip_extra",
|
|
"skip_forward",
|
|
"input_names_after_prune",
|
|
"skip_prune_program",
|
|
"separate_parameters",
|
|
]
|
|
|
|
# input check
|
|
for key in configs:
|
|
if key not in supported_configs:
|
|
raise ValueError(
|
|
f"The additional config ({key}) of `paddle.jit.save` is not supported."
|
|
)
|
|
|
|
# construct inner config
|
|
inner_config = _SaveLoadConfig()
|
|
inner_config.output_spec = configs.get("output_spec", None)
|
|
inner_config.with_hook = configs.get("with_hook", False)
|
|
inner_config.combine_params = configs.get("combine_params", False)
|
|
inner_config.clip_extra = configs.get("clip_extra", True)
|
|
inner_config.skip_forward = configs.get("skip_forward", False)
|
|
inner_config.input_names_after_prune = configs.get(
|
|
"input_names_after_prune", None
|
|
)
|
|
inner_config.skip_prune_program = configs.get("skip_prune_program", False)
|
|
inner_config.separate_parameters = configs.get("separate_parameters", False)
|
|
|
|
return inner_config
|
|
|
|
|
|
def _parse_load_config(configs: _LoadOptions) -> _SaveLoadConfig:
|
|
supported_configs = ['model_filename', 'params_filename']
|
|
|
|
# input check
|
|
for key in configs:
|
|
if key not in supported_configs:
|
|
raise ValueError(
|
|
f"The additional config ({key}) of `paddle.jit.load` is not supported."
|
|
)
|
|
|
|
# construct inner config
|
|
inner_config = _SaveLoadConfig()
|
|
inner_config.model_filename = configs.get('model_filename', None)
|
|
inner_config.params_filename = configs.get('params_filename', None)
|
|
|
|
return inner_config
|
|
|
|
|
|
def _get_input_var_and_names(inputs, input_spec, input_names_after_prune):
|
|
name_none_error = (
|
|
"The %s's name is None. "
|
|
"When using jit.save, please set InputSpec's name in "
|
|
"to_static(input_spec=[]) and jit.save(input_spec=[]) "
|
|
"and make sure they are consistent."
|
|
)
|
|
name_no_exists_error = (
|
|
"The tensor `%s` does not exists. "
|
|
"Please make sure the name of InputSpec or example Tensor "
|
|
"in input_spec is the same as the name of InputSpec in "
|
|
"`to_static` decorated on the Layer.forward method."
|
|
)
|
|
if input_names_after_prune is not None:
|
|
input_spec = [
|
|
x
|
|
for x in input_spec
|
|
if isinstance(x, paddle.static.InputSpec)
|
|
and x.name in input_names_after_prune
|
|
]
|
|
|
|
input_vars = [
|
|
var
|
|
for var in paddle.utils.flatten(inputs)
|
|
if isinstance(var, (Variable, paddle.pir.Value))
|
|
]
|
|
input_var_names = [
|
|
var.name
|
|
for var in paddle.utils.flatten(inputs)
|
|
if isinstance(var, (Variable, paddle.pir.Value))
|
|
]
|
|
|
|
if input_spec is None:
|
|
# no prune
|
|
return input_vars, input_var_names
|
|
else:
|
|
# filter out non-tensor type spec infos.
|
|
input_spec = [
|
|
spec
|
|
for spec in input_spec
|
|
if isinstance(spec, paddle.static.InputSpec)
|
|
]
|
|
result_var_list = []
|
|
result_name_list = []
|
|
if len(input_spec) == len(input_var_names):
|
|
# no prune
|
|
result_var_list = input_vars
|
|
result_name_list = input_var_names
|
|
# if input spec name not in input_var_names, only raise warning
|
|
for spec in input_spec:
|
|
if spec.name is None:
|
|
warnings.warn(name_none_error % spec)
|
|
elif spec.name not in input_var_names:
|
|
warnings.warn(name_no_exists_error % spec.name)
|
|
else:
|
|
pass
|
|
else:
|
|
# prune
|
|
for spec in input_spec:
|
|
if spec.name is None:
|
|
# name is None, the input_spec only can be InputSpec
|
|
raise ValueError(name_none_error % spec)
|
|
elif spec.name not in input_var_names:
|
|
# the input_spec can be `InputSpec` or `Tensor`
|
|
raise ValueError(name_no_exists_error % spec.name)
|
|
else:
|
|
result_var_list.append(spec)
|
|
result_name_list.append(spec.name)
|
|
|
|
return result_var_list, result_name_list
|
|
|
|
|
|
def _contains_dict(output):
|
|
if isinstance(output, dict):
|
|
return True
|
|
if isinstance(output, Sequence) and not isinstance(output, str):
|
|
return any(_contains_dict(i) for i in output)
|
|
return False
|
|
|
|
|
|
def _get_output_vars(outputs, output_spec, with_hook=False):
|
|
name_no_exists_error = (
|
|
"The tensor `%s` does not exists. "
|
|
"Please make sure the name of example Tensor "
|
|
"in configs.output_spec is the output tensor of "
|
|
"Layer.forward method."
|
|
)
|
|
output_spec_is_not_value_error = (
|
|
"tensor `%s` is not support in pir mode, "
|
|
"because pir value has no name sometimes, especially as output,"
|
|
"so we can't check tensor's name with output var name, please"
|
|
"change as pir.value(to_static layer's output)"
|
|
"or int(the position of to_static layer's output)"
|
|
)
|
|
if output_spec and with_hook:
|
|
raise RuntimeError(
|
|
"Currently not support specify output_spec while founding pre/post hooks in your outermost layer."
|
|
)
|
|
result_list = []
|
|
if _contains_dict(outputs):
|
|
warnings.warn(
|
|
"Found 'dict' in given outputs, the values will be returned in a sequence sorted in lexicographical order by their keys."
|
|
)
|
|
if use_pir_api():
|
|
from paddle.autograd.backward_utils import ValueSet
|
|
|
|
for var in paddle.utils.flatten(outputs):
|
|
if isinstance(var, paddle.pir.Value) and var not in ValueSet(
|
|
result_list
|
|
):
|
|
result_list.append(var)
|
|
|
|
if output_spec is not None:
|
|
output_size = len(result_list)
|
|
if len(output_spec) == output_size:
|
|
for var in output_spec:
|
|
if not isinstance(var, (paddle.pir.Value, int)):
|
|
warnings.warn(output_spec_is_not_value_error % var.name)
|
|
else:
|
|
if var not in ValueSet(result_list):
|
|
warnings.warn(name_no_exists_error % var.name)
|
|
else:
|
|
result_set = ValueSet(result_list)
|
|
part_result_list = []
|
|
for var in output_spec:
|
|
if isinstance(var, paddle.pir.Value):
|
|
if var not in result_set:
|
|
raise ValueError(name_no_exists_error % var.name)
|
|
else:
|
|
part_result_list.append(var)
|
|
elif isinstance(var, int):
|
|
if var >= output_size:
|
|
raise ValueError(
|
|
"position %d should smaller than output's size % d",
|
|
var,
|
|
output_size,
|
|
)
|
|
else:
|
|
part_result_list.append(result_list[var])
|
|
|
|
else:
|
|
raise ValueError(
|
|
output_spec_is_not_value_error % var.name
|
|
)
|
|
|
|
return part_result_list
|
|
else:
|
|
output_vars_dict = OrderedDict()
|
|
for var in paddle.utils.flatten(outputs):
|
|
if isinstance(var, (Variable)):
|
|
output_vars_dict[var.name] = var
|
|
if output_spec is None:
|
|
result_list = list(output_vars_dict.values())
|
|
elif output_spec is not None and len(output_spec) == len(
|
|
output_vars_dict
|
|
):
|
|
result_list = list(output_vars_dict.values())
|
|
for var in output_spec:
|
|
if var.name not in output_vars_dict:
|
|
warnings.warn(name_no_exists_error % var.name)
|
|
else:
|
|
for var in output_spec:
|
|
if var.name not in output_vars_dict:
|
|
raise ValueError(name_no_exists_error % var.name)
|
|
else:
|
|
result_list.append(output_vars_dict[var.name])
|
|
|
|
return result_list
|
|
|
|
|
|
# NOTE(chenweihang): [ Handling of use cases of API paddle.jit.load ]
|
|
# `paddle.jit.load` may be used to load saved results of:
|
|
# 1. Expected cases:
|
|
# - paddle.jit.save
|
|
# - paddle.static.save_inference_model
|
|
# 2. Error cases:
|
|
# - paddle.save: no .pdmodel for prefix
|
|
# - paddle.static.save: no .pdiparams but .pdparams exists
|
|
# - paddle.base.io.save_params/save_persistables: no __model__
|
|
# TODO(chenweihang): polish error message in above error cases
|
|
def _build_load_path_and_config(path, config):
|
|
# NOTE(chenweihang): If both [prefix save format] and [directory save format] exist,
|
|
# raise error, avoid confusing behavior
|
|
if use_pir_api():
|
|
model_suffix = PIR_INFER_MODEL_SUFFIX
|
|
else:
|
|
model_suffix = INFER_MODEL_SUFFIX
|
|
prefix_format_path = path + model_suffix
|
|
prefix_format_exist = os.path.exists(prefix_format_path)
|
|
directory_format_exist = os.path.isdir(path)
|
|
if prefix_format_exist and directory_format_exist:
|
|
raise ValueError(
|
|
f"The {path}.pdmodel(json) and {path} directory exist at the same time, "
|
|
"don't know which one to load, please make sure that the specified target "
|
|
"of ``path`` is unique."
|
|
)
|
|
elif not prefix_format_exist and not directory_format_exist:
|
|
raise ValueError(
|
|
f"The ``path`` ({path}) to load model not exists. "
|
|
"Please make sure that *.pdmodel(json) exists or "
|
|
"don't using ``skip_forward=True`` to jit.save."
|
|
)
|
|
else:
|
|
if prefix_format_exist:
|
|
file_prefix = os.path.basename(path)
|
|
model_path = os.path.dirname(path)
|
|
if config.model_filename is not None:
|
|
warnings.warn(
|
|
"When loading the result saved with the "
|
|
"specified file prefix, the ``model_filename`` config does "
|
|
"not take effect."
|
|
)
|
|
config.model_filename = file_prefix + model_suffix
|
|
if config.params_filename is not None:
|
|
warnings.warn(
|
|
"When loading the result saved with the "
|
|
"specified file prefix, the ``params_filename`` config does "
|
|
"not take effect."
|
|
)
|
|
config.params_filename = file_prefix + INFER_PARAMS_SUFFIX
|
|
else:
|
|
# Compatible with the old save_inference_model format
|
|
model_path = path
|
|
|
|
return model_path, config
|
|
|
|
|
|
_save_pre_hooks_lock = threading.Lock()
|
|
_save_pre_hooks = []
|
|
|
|
|
|
class HookRemoveHelper:
|
|
"""A HookRemoveHelper that can be used to remove hook."""
|
|
|
|
def __init__(self, hook):
|
|
self._hook = hook
|
|
|
|
def remove(self):
|
|
_remove_save_pre_hook(self._hook)
|
|
|
|
|
|
def _register_save_pre_hook(hook):
|
|
"""
|
|
Register a save pre-hook for `paddle.jit.save`.
|
|
This hook will be executed before `save` function has been invoked.
|
|
|
|
hook(layer, input_spec, configs) -> None
|
|
- layer (Layer|function): This argument is corresponding to `layer` in `paddle.jit.save`.
|
|
- input_spec (list or tuple[InputSpec|Tensor|Python built-in variable]): This argument is corresponding to `input_spec` in `paddle.jit.save`.
|
|
- configs (dict): This argument is corresponding to `configs` in `paddle.jit.save`.
|
|
|
|
Args:
|
|
hook(function): a function registered as a save pre-hook
|
|
|
|
Returns:
|
|
HookRemoveHelper: a HookRemoveHelper object that can be used to remove the added hook by calling `hook_remove_helper.remove()`.
|
|
|
|
Examples:
|
|
.. code-block:: pycon
|
|
|
|
>>> # doctest: +SKIP('`paddle.jit.api.to_static` can not run in xdoctest')
|
|
>>> import numpy as np
|
|
>>> import paddle
|
|
|
|
>>> IMAGE_SIZE = 256
|
|
>>> CLASS_NUM = 10
|
|
|
|
>>> class LinearNet(paddle.nn.Layer):
|
|
... def __init__(self):
|
|
... super().__init__()
|
|
... self._linear = paddle.nn.Linear(IMAGE_SIZE, CLASS_NUM)
|
|
...
|
|
... def forward(self, x):
|
|
... return self._linear(x)
|
|
>>> saving_count = 0
|
|
>>> def save_pre_hook(layer, input_spec, configs):
|
|
... global saving_count
|
|
... saving_count += 1
|
|
>>> remove_handler = paddle.jit.api._register_save_pre_hook(save_pre_hook)
|
|
|
|
>>> layer = LinearNet()
|
|
>>> paddle.jit.save(layer, "/tmp", [paddle.static.InputSpec(shape=[-1, IMAGE_SIZE])])
|
|
>>> print(saving_count)
|
|
1
|
|
|
|
>>> remove_handler.remove()
|
|
>>> paddle.jit.save(layer, "/tmp", [paddle.static.InputSpec(shape=[-1, IMAGE_SIZE])])
|
|
>>> print(saving_count)
|
|
1
|
|
"""
|
|
global _save_pre_hooks_lock
|
|
global _save_pre_hooks
|
|
_save_pre_hooks_lock.acquire()
|
|
if hook not in _save_pre_hooks:
|
|
_save_pre_hooks.append(hook)
|
|
_save_pre_hooks_lock.release()
|
|
return HookRemoveHelper(hook)
|
|
|
|
|
|
def _clear_save_pre_hooks():
|
|
global _save_pre_hooks_lock
|
|
global _save_pre_hooks
|
|
_save_pre_hooks_lock.acquire()
|
|
_save_pre_hooks.clear()
|
|
_save_pre_hooks_lock.release()
|
|
|
|
|
|
def _remove_save_pre_hook(hook):
|
|
global _save_pre_hooks_lock
|
|
global _save_pre_hooks
|
|
_save_pre_hooks_lock.acquire()
|
|
if hook in _save_pre_hooks:
|
|
_save_pre_hooks.remove(hook)
|
|
_save_pre_hooks_lock.release()
|
|
|
|
|
|
class _SaveFunction(Protocol):
|
|
def __call__(
|
|
self,
|
|
layer: Layer | Callable[..., Any],
|
|
path: str,
|
|
input_spec: Sequence[InputSpec | Tensor | object] | None = ...,
|
|
**configs: Unpack[_SaveOptions],
|
|
) -> None: ...
|
|
|
|
|
|
@wrap_decorator
|
|
def _run_save_pre_hooks(func: _SaveFunction) -> _SaveFunction:
|
|
def wrapper(
|
|
layer: Layer | Callable[..., Any],
|
|
path: str,
|
|
input_spec: Sequence[InputSpec | Tensor | object] | None = None,
|
|
**configs: Unpack[_SaveOptions],
|
|
) -> None:
|
|
global _save_pre_hooks
|
|
for hook in _save_pre_hooks:
|
|
hook(layer, input_spec, configs)
|
|
func(layer, path, input_spec, **configs)
|
|
|
|
return wrapper
|
|
|
|
|
|
def _save_property(filename: str, property_vals: list[tuple[Any, str]]):
|
|
"""class property serialization.
|
|
|
|
Args:
|
|
filename (str): *.meta
|
|
property_vals (list[tuple[Any, str]]): class property.
|
|
"""
|
|
|
|
def set_property(meta, key, val):
|
|
if isinstance(val, float):
|
|
meta.set_float(key, val)
|
|
elif isinstance(val, int):
|
|
meta.set_int(key, val)
|
|
elif isinstance(val, str):
|
|
meta.set_string(key, val)
|
|
elif isinstance(val, (tuple, list)):
|
|
if isinstance(val[0], float):
|
|
meta.set_floats(key, val)
|
|
elif isinstance(val[0], int):
|
|
meta.set_ints(key, val)
|
|
elif isinstance(val[0], str):
|
|
meta.set_strings(key, val)
|
|
else:
|
|
raise ValueError(f"Note support val type: {type(val)}")
|
|
|
|
with open(filename, 'wb') as f:
|
|
meta = paddle.framework.core.Property()
|
|
for item in property_vals:
|
|
val, key = item[0], item[1]
|
|
set_property(meta, key, val)
|
|
f.write(meta.serialize_to_string())
|
|
|
|
|
|
def _get_function_names_from_layer(layer: Layer) -> list[str]:
|
|
cls = layer.__class__
|
|
return [
|
|
member_name
|
|
for member_name, member in inspect.getmembers(cls)
|
|
if (
|
|
inspect.isfunction(member)
|
|
or inspect.ismethod(member)
|
|
or inspect.ismethoddescriptor(member)
|
|
)
|
|
]
|
|
|
|
|
|
@_run_save_pre_hooks
|
|
@switch_to_static_graph
|
|
def save(
|
|
layer: Layer | Callable[..., Any],
|
|
path: str,
|
|
input_spec: Sequence[InputSpec | Tensor | object] | None = None,
|
|
**configs: Unpack[_SaveOptions],
|
|
) -> None:
|
|
"""
|
|
Saves input Layer or function as ``paddle.jit.TranslatedLayer``
|
|
format model, which can be used for inference or fine-tuning after loading.
|
|
|
|
It will save the translated program and all related persistable
|
|
variables of input Layer to given ``path`` .
|
|
|
|
``path`` is the prefix of saved objects, and the saved translated program file
|
|
suffix is ``.pdmodel`` , the saved persistable variables file suffix is ``.pdiparams`` ,
|
|
and here also saved some additional variable description information to a file,
|
|
its suffix is ``.pdiparams.info``, these additional information is used in fine-tuning.
|
|
|
|
The saved model can be loaded by follow APIs:
|
|
- ``paddle.jit.load``
|
|
- ``paddle.static.load_inference_model``
|
|
- Other C++ inference APIs
|
|
|
|
.. note::
|
|
When using ``paddle.jit.save`` to save a function, parameters will not be saved. If you have to
|
|
save the parameter, please pass the Layer containing function and parameter to ``paddle.jit.save``.
|
|
|
|
Args:
|
|
layer (Layer|function): The Layer or function to be saved.
|
|
path (str): The path prefix to save model. The format is ``dirname/file_prefix`` or ``file_prefix``.
|
|
input_spec (list or tuple[InputSpec|Tensor|Python built-in variable], optional): Describes the input of the saved model's forward
|
|
method, which can be described by InputSpec or example Tensor. Moreover, we support to specify non-tensor type argument,
|
|
such as int, float, string, or list/dict of them.If None, all input variables of
|
|
the original Layer's forward method would be the inputs of the saved model. Default None.
|
|
**configs (dict, optional): Other save configuration options for compatibility. We do not
|
|
recommend using these configurations, they may be removed in the future. If not necessary,
|
|
DO NOT use them. Default None.
|
|
The following options are currently supported:
|
|
(1) output_spec (list[Tensor|Value|int]): Selects the output targets of the saved model,
|
|
By default, all return variables of original Layer's forward method are kept as the
|
|
output of the saved model. If the provided ``output_spec`` list is not all output variables,
|
|
the saved model will be pruned according to the given ``output_spec`` list.
|
|
in pir mode, Tensor is not supported, because value has no name in most cases,
|
|
which can't be used to judge which tensor corresponds to which value; the value can't be found
|
|
if the saved program is not the same as the program that includes output_spec, so we need to
|
|
use the position of the output.
|
|
|
|
Returns:
|
|
None
|
|
|
|
Examples:
|
|
.. code-block:: pycon
|
|
|
|
>>> # doctest: +SKIP('`paddle.jit.to_static` can not run in xdoctest')
|
|
>>> # example 1: save layer
|
|
>>> import numpy as np
|
|
>>> import paddle
|
|
>>> import paddle.nn as nn
|
|
>>> import paddle.optimizer as opt
|
|
|
|
>>> BATCH_SIZE = 16
|
|
>>> BATCH_NUM = 4
|
|
>>> EPOCH_NUM = 4
|
|
|
|
>>> IMAGE_SIZE = 784
|
|
>>> CLASS_NUM = 10
|
|
|
|
>>> # define a random dataset
|
|
>>> class RandomDataset(paddle.io.Dataset): # type: ignore[type-arg]
|
|
... def __init__(self, num_samples):
|
|
... self.num_samples = num_samples
|
|
...
|
|
... def __getitem__(self, idx):
|
|
... image = np.random.random([IMAGE_SIZE]).astype('float32')
|
|
... label = np.random.randint(0, CLASS_NUM - 1, (1,)).astype('int64')
|
|
... return image, label
|
|
...
|
|
... def __len__(self):
|
|
... return self.num_samples
|
|
|
|
>>> class LinearNet(nn.Layer):
|
|
... def __init__(self):
|
|
... super().__init__()
|
|
... self._linear = nn.Linear(IMAGE_SIZE, CLASS_NUM)
|
|
...
|
|
... @paddle.jit.to_static
|
|
... def forward(self, x):
|
|
... return self._linear(x)
|
|
|
|
>>> def train(layer, loader, loss_fn, opt):
|
|
... for epoch_id in range(EPOCH_NUM):
|
|
... for batch_id, (image, label) in enumerate(loader()):
|
|
... out = layer(image)
|
|
... loss = loss_fn(out, label)
|
|
... loss.backward()
|
|
... opt.step()
|
|
... opt.clear_grad()
|
|
... print("Epoch {} batch {}: loss = {}".format(epoch_id, batch_id, np.mean(loss.numpy())))
|
|
|
|
>>> # 1. train & save model.
|
|
|
|
>>> # create network
|
|
>>> layer = LinearNet()
|
|
>>> loss_fn = nn.CrossEntropyLoss()
|
|
>>> adam = opt.Adam(learning_rate=0.001, parameters=layer.parameters())
|
|
|
|
>>> # create data loader
|
|
>>> dataset = RandomDataset(BATCH_NUM * BATCH_SIZE)
|
|
>>> loader = paddle.io.DataLoader(
|
|
... dataset,
|
|
... batch_size=BATCH_SIZE,
|
|
... shuffle=True,
|
|
... drop_last=True,
|
|
... num_workers=2,
|
|
... )
|
|
|
|
>>> # train
|
|
>>> train(layer, loader, loss_fn, adam)
|
|
|
|
>>> # save
|
|
>>> path = "example_model/linear"
|
|
>>> paddle.jit.save(layer, path)
|
|
|
|
>>> # example 2: save function
|
|
>>> import paddle
|
|
>>> from paddle.static import InputSpec
|
|
|
|
|
|
>>> def save_function():
|
|
... @paddle.jit.to_static
|
|
... def fun(inputs):
|
|
... return paddle.tanh(inputs)
|
|
...
|
|
... path = 'test_jit_save_load_function_1/func'
|
|
... inps = paddle.rand([3, 6])
|
|
... origin = fun(inps)
|
|
...
|
|
... paddle.jit.save(fun, path)
|
|
... load_func = paddle.jit.load(path)
|
|
...
|
|
... load_result = load_func(inps)
|
|
... print((load_result - origin).abs().max() < 1e-10)
|
|
|
|
>>> save_function()
|
|
"""
|
|
# 1. input build & check
|
|
prog_translator = ProgramTranslator()
|
|
is_prim_infer = core._is_fwd_prim_enabled() and core._is_bwd_prim_enabled()
|
|
if not prog_translator.enable_to_static:
|
|
raise RuntimeError(
|
|
"The paddle.jit.save doesn't work when setting 'paddle.jit.enable_to_static' to False."
|
|
)
|
|
|
|
if not (
|
|
isinstance(layer, (Layer, StaticFunction)) or inspect.isfunction(layer)
|
|
):
|
|
raise TypeError(
|
|
f"The input of paddle.jit.save should be 'Layer' or 'Function', but received input type is {type(layer)}."
|
|
)
|
|
elif inspect.isfunction(layer) or isinstance(layer, StaticFunction):
|
|
warnings.warn(
|
|
'What you save is a function, and `jit.save` will generate the name of the model file according to `path` you specify. When loading these files with `jit.load`, you get a `TranslatedLayer` whose inference result is the same as the inference result of the function you saved.'
|
|
)
|
|
|
|
# NOTE(chenweihang): If the input layer be wrapped by DataParallel,
|
|
# the args and kwargs of forward method will can't be parsed by
|
|
# function_spec, so here we save DataParallel._layers instead
|
|
# DataParallel it self
|
|
# NOTE(chenweihang): using inner_layer, do not change input layer
|
|
if isinstance(layer, paddle.DataParallel):
|
|
inner_layer = layer._layers
|
|
else:
|
|
inner_layer = layer
|
|
|
|
# path check
|
|
file_prefix = os.path.basename(path)
|
|
if file_prefix == "":
|
|
raise ValueError(
|
|
"The input path MUST be format of dirname/file_prefix "
|
|
"[dirname\\file_prefix in Windows system], but received "
|
|
"file_prefix is empty string."
|
|
)
|
|
|
|
dirname = os.path.dirname(path)
|
|
if dirname and not os.path.exists(dirname):
|
|
os.makedirs(dirname)
|
|
|
|
# avoid change user given input_spec
|
|
inner_input_spec = None
|
|
if input_spec is not None:
|
|
if isinstance(layer, Layer):
|
|
for member_name in _get_function_names_from_layer(inner_layer):
|
|
static_func = getattr(inner_layer, member_name, None)
|
|
if (
|
|
isinstance(static_func, StaticFunction)
|
|
and 'forward' != member_name
|
|
):
|
|
raise ValueError(
|
|
f"If there are static functions other than 'forward' that need to be saved, the input 'input_spec' should be None, but received the type of 'input_spec' is {type(input_spec)}."
|
|
)
|
|
|
|
if not isinstance(input_spec, (list, tuple)):
|
|
raise TypeError(
|
|
f"The input input_spec should be 'list', but received input_spec's type is {type(input_spec)}."
|
|
)
|
|
inner_input_spec = []
|
|
for var in paddle.utils.flatten(input_spec):
|
|
if isinstance(var, paddle.static.InputSpec):
|
|
inner_input_spec.append(var)
|
|
elif isinstance(
|
|
var, (core.eager.Tensor, Variable, paddle.pir.Value)
|
|
):
|
|
inner_input_spec.append(
|
|
paddle.static.InputSpec.from_tensor(var)
|
|
)
|
|
else:
|
|
# NOTE(Aurelius84): Support non-Tensor type in `input_spec`.
|
|
inner_input_spec.append(var)
|
|
|
|
# parse configs
|
|
configs = _parse_save_configs(configs)
|
|
# whether outermost layer has pre/post hook, if does, we need also save
|
|
# these operators in program.
|
|
with_hook = configs.with_hook
|
|
combine_params = configs.combine_params
|
|
if combine_params:
|
|
configs._program_only = True
|
|
|
|
scope = core.Scope()
|
|
extra_var_info = {}
|
|
if isinstance(layer, Layer):
|
|
functions = list(set(_get_function_names_from_layer(inner_layer)))
|
|
functions = sorted(functions)
|
|
if inner_layer._forward_pre_hooks or inner_layer._forward_post_hooks:
|
|
with_hook = True
|
|
else:
|
|
# layer is function
|
|
functions = [layer]
|
|
|
|
combine_vars = {}
|
|
combine_program = []
|
|
property_vals = [] # (value, key)
|
|
concrete_program = None
|
|
for attr_func in functions:
|
|
if isinstance(layer, Layer):
|
|
static_func = get_ast_static_function(
|
|
getattr(inner_layer, attr_func, None)
|
|
)
|
|
if isinstance(static_func, StaticFunction):
|
|
if static_func.is_property:
|
|
# property method to be exported
|
|
immediate_val = static_func()
|
|
property_vals.append(
|
|
(
|
|
immediate_val,
|
|
layer.__class__.__name__ + '.' + attr_func,
|
|
)
|
|
)
|
|
continue
|
|
|
|
concrete_program = (
|
|
static_func.concrete_program_specify_input_spec(
|
|
inner_input_spec,
|
|
with_hook=with_hook,
|
|
is_prim_infer=is_prim_infer,
|
|
)
|
|
)
|
|
|
|
elif 'forward' == attr_func:
|
|
if configs.skip_forward:
|
|
# do not jit.save forward function
|
|
continue
|
|
|
|
# transform in jit.save, if input_spec is incomplete, declarative will throw error
|
|
# inner_input_spec is list[InputSpec], it should be packed with same structure
|
|
# as original input_spec here.
|
|
if inner_input_spec:
|
|
inner_input_spec = paddle.utils.pack_sequence_as(
|
|
input_spec, inner_input_spec
|
|
)
|
|
static_forward = to_static(
|
|
inner_layer.forward,
|
|
input_spec=inner_input_spec,
|
|
full_graph=True,
|
|
backend=None,
|
|
)
|
|
|
|
concrete_program = (
|
|
static_forward.concrete_program_specify_input_spec(
|
|
with_hook=with_hook, is_prim_infer=is_prim_infer
|
|
)
|
|
)
|
|
# the input_spec has been used in declarative, which is equal to
|
|
# @to_static with input_spec and jit.save without input_spec,
|
|
# avoid needless warning
|
|
inner_input_spec = None
|
|
else:
|
|
continue
|
|
else:
|
|
# When layer is a function
|
|
if isinstance(attr_func, StaticFunction):
|
|
static_func = get_ast_static_function(attr_func)
|
|
|
|
if static_func.is_property:
|
|
# property method to be exported
|
|
immediate_val = static_func()
|
|
property_vals.append((immediate_val, static_func))
|
|
continue
|
|
|
|
concrete_program = (
|
|
static_func.concrete_program_specify_input_spec(
|
|
inner_input_spec, is_prim_infer=is_prim_infer
|
|
)
|
|
)
|
|
else:
|
|
static_func = get_ast_static_function(attr_func)
|
|
if inner_input_spec:
|
|
inner_input_spec = paddle.utils.pack_sequence_as(
|
|
input_spec, inner_input_spec
|
|
)
|
|
static_function = to_static(
|
|
static_func,
|
|
input_spec=inner_input_spec,
|
|
full_graph=True,
|
|
backend=None,
|
|
)
|
|
concrete_program = static_function.concrete_program
|
|
|
|
if static_function.class_instance is None:
|
|
warnings.warn(
|
|
f'`jit.save` will only save the `Program`, not the parameters. If you have to save the parameters, please make sure that {layer} is a member function of `paddle.nn.Layer` and the saved parameters are in `state_dict`'
|
|
)
|
|
|
|
# when save multi `StaticFunction`, all `StaticFunction` share params.
|
|
dygraph_state_dict = None
|
|
if isinstance(inner_layer, Layer):
|
|
dygraph_state_dict = inner_layer.to_static_state_dict()
|
|
elif isinstance(attr_func, StaticFunction):
|
|
if static_func.class_instance:
|
|
dygraph_state_dict = (
|
|
static_func.class_instance.to_static_state_dict()
|
|
)
|
|
|
|
if dygraph_state_dict:
|
|
# NOTE(chenweihang): we maintain the mapping of variable name to
|
|
# structured name, the buffer variable (non-persistable)
|
|
# saved to inference program may not need by dygraph Layer,
|
|
# we only record the state_dict variable's structured name
|
|
state_names_dict = {}
|
|
state_var_dict = {}
|
|
for structured_name, var in dygraph_state_dict.items():
|
|
state_names_dict[var.name] = structured_name
|
|
state_var_dict[var.name] = var
|
|
# 3. share parameters from Layer to scope & record var info
|
|
with dygraph.guard():
|
|
if use_pir_api():
|
|
for tensor, value in zip(*concrete_program.parameters):
|
|
if not value.persistable:
|
|
continue
|
|
param_or_buffer_tensor = scope.var(value.name).get_tensor()
|
|
src_tensor = (
|
|
state_var_dict[tensor.name].value().get_tensor()
|
|
)
|
|
param_or_buffer_tensor._share_data_with(src_tensor)
|
|
|
|
else:
|
|
for param_or_buffer in concrete_program.parameters:
|
|
# share to scope
|
|
if param_or_buffer.type == core.VarDesc.VarType.VOCAB:
|
|
scr_tensor = param_or_buffer.value().get_map_tensor()
|
|
tgt_var = scope.var(param_or_buffer.name)
|
|
tgt_var.set_vocab(scr_tensor)
|
|
else:
|
|
param_or_buffer_tensor = scope.var(
|
|
param_or_buffer.name
|
|
).get_tensor()
|
|
# src_tensor = param_or_buffer.value().get_tensor()
|
|
src_tensor = (
|
|
state_var_dict[param_or_buffer.name]
|
|
.value()
|
|
.get_tensor()
|
|
)
|
|
param_or_buffer_tensor._share_data_with(src_tensor)
|
|
# record var info
|
|
if param_or_buffer.name not in extra_var_info:
|
|
extra_info_dict = {}
|
|
if param_or_buffer.name in state_names_dict:
|
|
extra_info_dict['structured_name'] = (
|
|
state_names_dict[param_or_buffer.name]
|
|
)
|
|
extra_info_dict['stop_gradient'] = (
|
|
param_or_buffer.stop_gradient
|
|
)
|
|
if isinstance(param_or_buffer, EagerParamBase):
|
|
extra_info_dict['trainable'] = (
|
|
param_or_buffer.trainable
|
|
)
|
|
extra_var_info[param_or_buffer.name] = extra_info_dict
|
|
# 4. build input & output of save_inference_model
|
|
# NOTE(chenweihang): [ Get input variables name ]
|
|
# There are two cases, whether to prune the inputs or not
|
|
# - not prune inputs (recommend):
|
|
# - the len(input_spec) == len((concrete_program.inputs) - 1
|
|
# - here can use concrete_program.inputs directly
|
|
# - prune inputs:
|
|
# - the input_spec length < len((concrete_program.inputs) - 1
|
|
# - the input_spec's name should be in concrete_program.inputs
|
|
|
|
input_vars, input_var_names = _get_input_var_and_names(
|
|
concrete_program.inputs,
|
|
inner_input_spec,
|
|
configs.input_names_after_prune,
|
|
)
|
|
|
|
# NOTE(chenweihang): [ Get output variables ]
|
|
# the rule is like [ Get input variables name ]. For output var,
|
|
# we only support Tensor spec, and actually, we only need the
|
|
# var name of output, and we don't recommended to use output_spec
|
|
# NOTE(Ruting): in pir mode, Tensor is not supported, because value has no name in most cases,
|
|
# which can't be used to judge which tensor corresponds to which value; the value can't be found
|
|
# if the saved program is not the same as the program that includes output_spec, so we need to
|
|
# use the position of the output.
|
|
|
|
output_vars = _get_output_vars(
|
|
concrete_program.outputs, configs.output_spec, with_hook
|
|
)
|
|
|
|
# 5. save inference model
|
|
# construct new save_inference_model arguments
|
|
model_path = dirname
|
|
# NOTE(chenweihang): because prefix contains model and params filename,
|
|
# so we don't support set model_filename & params_filename
|
|
if 'forward' == attr_func or not isinstance(layer, Layer):
|
|
model_filename = file_prefix + INFER_MODEL_SUFFIX
|
|
params_filename = file_prefix + INFER_PARAMS_SUFFIX
|
|
path_prefix = file_prefix
|
|
else:
|
|
model_filename = file_prefix + '.' + attr_func + INFER_MODEL_SUFFIX
|
|
params_filename = (
|
|
file_prefix + '.' + attr_func + INFER_PARAMS_SUFFIX
|
|
)
|
|
path_prefix = file_prefix + '.' + attr_func
|
|
file_path = os.path.join(model_path, path_prefix)
|
|
with scope_guard(scope):
|
|
if use_pir_api():
|
|
value_map = paddle.pir.IrMapping()
|
|
clone_program = concrete_program.main_program.clone(value_map)
|
|
clone_input_vars = []
|
|
for v in input_vars:
|
|
if type(v) is paddle.static.InputSpec:
|
|
name = v.name
|
|
for op in clone_program.global_block().ops:
|
|
if (
|
|
op.name() == 'pd_op.data'
|
|
and op.attrs()["name"] == name
|
|
):
|
|
clone_input_vars.append(op.result(0))
|
|
else:
|
|
clone_input_vars.append(value_map.look_up(v))
|
|
|
|
clone_output_vars = [value_map.look_up(v) for v in output_vars]
|
|
|
|
else:
|
|
input_vars = [
|
|
concrete_program.main_program.global_block().var(name)
|
|
for name in input_var_names
|
|
]
|
|
clone_program = concrete_program.main_program.clone()
|
|
clone_input_vars = input_vars
|
|
clone_output_vars = output_vars
|
|
save_inference_model(
|
|
path_prefix=file_path,
|
|
feed_vars=clone_input_vars,
|
|
fetch_vars=clone_output_vars,
|
|
executor=Executor(_current_expected_place()),
|
|
program=clone_program,
|
|
clip_extra=configs.clip_extra,
|
|
skip_prune_program=configs.skip_prune_program,
|
|
separate_parameters=configs.separate_parameters,
|
|
)
|
|
|
|
if combine_params:
|
|
if use_pir_api():
|
|
# NOTE(Ruting): concrete_program has been pruned when init partialProgramLayer,
|
|
# so we do not need to prune again.
|
|
|
|
for var in concrete_program.main_program.list_vars():
|
|
if var.persistable:
|
|
combine_vars[var.name] = var
|
|
# NOTE(Ruting): concrete_program will delete after this loop item,
|
|
# value delete at the same time, so we use list to Extend its lifecycle
|
|
combine_program.append(concrete_program.main_program)
|
|
else:
|
|
clone_main_program = concrete_program.main_program.clone()
|
|
clone_main_program = clone_main_program._prune_with_input(
|
|
input_var_names, output_vars
|
|
)
|
|
for block in clone_main_program.blocks:
|
|
combine_vars.update(block.vars)
|
|
|
|
# save shared params
|
|
if combine_params:
|
|
# sort vars by name
|
|
combine_vars = sorted(combine_vars.items(), key=lambda item: item[0])
|
|
ordered_vars = []
|
|
for name, var in combine_vars:
|
|
ordered_vars.append(var)
|
|
|
|
params_filename = file_prefix + INFER_PARAMS_SUFFIX
|
|
with scope_guard(scope):
|
|
if use_pir_api():
|
|
paddle.static.save_vars(
|
|
Executor(_current_expected_place()),
|
|
dirname=model_path,
|
|
vars=ordered_vars,
|
|
filename=params_filename,
|
|
)
|
|
else:
|
|
paddle.static.save_vars(
|
|
Executor(_current_expected_place()),
|
|
dirname=model_path,
|
|
vars=list(
|
|
filter(
|
|
paddle.framework.io_utils.is_persistable,
|
|
ordered_vars,
|
|
)
|
|
),
|
|
filename=params_filename,
|
|
)
|
|
# save property
|
|
property_save_path = os.path.join(
|
|
os.path.normpath(model_path), file_prefix + INFER_PROPERTY_SUFFIX
|
|
)
|
|
_save_property(property_save_path, property_vals)
|
|
|
|
# NOTE(chenweihang): [ Save extra variable info ]
|
|
# save_inference_model will lose some important variable information, including:
|
|
# - Variable name and correspondence (when saved variables as one file)
|
|
# - Variable.stop_gradient information
|
|
# - Which persistent variable are parameter and which are not
|
|
# - Parameter.trainable information
|
|
#
|
|
# The lost information cannot be recovered when it is loaded again,
|
|
# so if we want to perform fine-tune after loading, we may need to
|
|
# configure redundant information to proceed.
|
|
#
|
|
# Due to compatibility issues, we cannot change the original storage structure,
|
|
# but we can save these information in `jit.save` without changing the original
|
|
# storage to improve user experience. So we save extra information into
|
|
# file `***.pdiparams.info`
|
|
|
|
# "layer" can only be Layer or function or StaticFunction.
|
|
contain_parameter = False
|
|
if concrete_program is not None:
|
|
for var in concrete_program.main_program.list_vars():
|
|
if use_pir_api():
|
|
is_persistable = (
|
|
var.get_defining_op().has_attr("persistable")
|
|
and var.get_defining_op().attrs()["persistable"] is True
|
|
)
|
|
contain_parameter |= is_persistable
|
|
else:
|
|
contain_parameter |= isinstance(var, Parameter)
|
|
|
|
if (isinstance(layer, Layer) or contain_parameter) and extra_var_info:
|
|
with scope_guard(scope):
|
|
extra_var_info_path = path + INFER_PARAMS_INFO_SUFFIX
|
|
with open(extra_var_info_path, 'wb') as f:
|
|
pickle.dump(extra_var_info, f, protocol=2)
|
|
scope.erase(scope.local_var_names())
|
|
|
|
|
|
@dygraph_only
|
|
def load(
|
|
path: str, **configs: Unpack[_LoadOptions]
|
|
) -> TranslatedLayer | PirTranslatedLayer:
|
|
"""
|
|
:api_attr: imperative
|
|
|
|
Load model saved by ``paddle.jit.save`` or ``paddle.static.save_inference_model`` or
|
|
paddle 1.x API ``paddle.static.save_inference_model`` as ``paddle.jit.TranslatedLayer``,
|
|
then performing inference or fine-tune training.
|
|
|
|
.. note::
|
|
If you load model saved by ``paddle.static.save_inference_model`` ,
|
|
there will be the following limitations when using it in fine-tuning:
|
|
1. Imperative mode do not support DenseTensor. All original model's feed targets or parameters that depend on LoD are temporarily unavailable.
|
|
2. All saved model's feed targets need to be passed into TranslatedLayer's forward function.
|
|
3. The variable's ``stop_gradient`` information is lost and can not be recovered.
|
|
4. The parameter's ``trainable`` information is lost and can not be recovered.
|
|
|
|
Args:
|
|
path (str): The path prefix to load model. The format is ``dirname/file_prefix`` or ``file_prefix`` .
|
|
**configs (dict, optional): Other load configuration options for compatibility. We do not
|
|
recommend using these configurations, they may be removed in the future. If not necessary,
|
|
DO NOT use them. Default None.
|
|
The following options are currently supported:
|
|
(1) model_filename (str): The inference model file name of the paddle 1.x
|
|
``save_inference_model`` save format. Default file name is :code:`__model__` .
|
|
(2) params_filename (str): The persistable variables file name of the paddle 1.x
|
|
``save_inference_model`` save format. No default file name, save variables separately
|
|
by default.
|
|
|
|
|
|
Returns:
|
|
TranslatedLayer: A Layer object can run saved translated model.
|
|
|
|
Examples:
|
|
1. Load model saved by ``paddle.jit.save`` then performing inference and fine-tune training.
|
|
|
|
.. code-block:: pycon
|
|
:name: code-example1
|
|
|
|
>>> # doctest: +SKIP('`paddle.jit.to_static` can not run in xdoctest')
|
|
>>> import numpy as np
|
|
>>> import paddle
|
|
>>> import paddle.nn as nn
|
|
>>> import paddle.optimizer as opt
|
|
|
|
>>> BATCH_SIZE = 16
|
|
>>> BATCH_NUM = 4
|
|
>>> EPOCH_NUM = 4
|
|
|
|
>>> IMAGE_SIZE = 784
|
|
>>> CLASS_NUM = 10
|
|
|
|
>>> # define a random dataset
|
|
>>> class RandomDataset(paddle.io.Dataset): # type: ignore[type-arg]
|
|
... def __init__(self, num_samples):
|
|
... self.num_samples = num_samples
|
|
...
|
|
... def __getitem__(self, idx):
|
|
... image = np.random.random([IMAGE_SIZE]).astype('float32')
|
|
... label = np.random.randint(0, CLASS_NUM - 1, (1,)).astype('int64')
|
|
... return image, label
|
|
...
|
|
... def __len__(self):
|
|
... return self.num_samples
|
|
|
|
>>> class LinearNet(nn.Layer):
|
|
... def __init__(self):
|
|
... super().__init__()
|
|
... self._linear = nn.Linear(IMAGE_SIZE, CLASS_NUM)
|
|
...
|
|
... @paddle.jit.to_static
|
|
... def forward(self, x):
|
|
... return self._linear(x)
|
|
>>> def train(layer, loader, loss_fn, opt):
|
|
... for epoch_id in range(EPOCH_NUM):
|
|
... for batch_id, (image, label) in enumerate(loader()):
|
|
... out = layer(image)
|
|
... loss = loss_fn(out, label)
|
|
... loss.backward()
|
|
... opt.step()
|
|
... opt.clear_grad()
|
|
... print("Epoch {} batch {}: loss = {}".format(epoch_id, batch_id, np.mean(loss.numpy())))
|
|
|
|
>>> # 1. train & save model.
|
|
|
|
>>> # create network
|
|
>>> layer = LinearNet()
|
|
>>> loss_fn = nn.CrossEntropyLoss()
|
|
>>> adam = opt.Adam(learning_rate=0.001, parameters=layer.parameters())
|
|
|
|
>>> # create data loader
|
|
>>> dataset = RandomDataset(BATCH_NUM * BATCH_SIZE)
|
|
>>> loader = paddle.io.DataLoader(
|
|
... dataset,
|
|
... batch_size=BATCH_SIZE,
|
|
... shuffle=True,
|
|
... drop_last=True,
|
|
... num_workers=2,
|
|
... )
|
|
|
|
>>> # train
|
|
>>> train(layer, loader, loss_fn, adam)
|
|
|
|
>>> # save
|
|
>>> path = "example_model/linear"
|
|
>>> paddle.jit.save(layer, path)
|
|
|
|
>>> # 2. load model
|
|
|
|
>>> # load
|
|
>>> loaded_layer = paddle.jit.load(path)
|
|
|
|
>>> # inference
|
|
>>> loaded_layer.eval()
|
|
>>> x = paddle.randn([1, IMAGE_SIZE], 'float32')
|
|
>>> pred = loaded_layer(x)
|
|
|
|
>>> # fine-tune
|
|
>>> loaded_layer.train()
|
|
>>> adam = opt.Adam(learning_rate=0.001, parameters=loaded_layer.parameters())
|
|
>>> train(loaded_layer, loader, loss_fn, adam)
|
|
|
|
|
|
2. Load model saved by ``paddle.static.save_inference_model`` then performing and fine-tune training.
|
|
|
|
.. code-block:: pycon
|
|
:name: code-example2
|
|
|
|
>>> # doctest: +SOLO('can not use multiprocessing testing `DataLoader`')
|
|
>>> import numpy as np
|
|
>>> import paddle
|
|
>>> import paddle.static as static
|
|
>>> import paddle.nn as nn
|
|
>>> import paddle.optimizer as opt
|
|
>>> import paddle.nn.functional as F
|
|
|
|
>>> BATCH_SIZE = 16
|
|
>>> BATCH_NUM = 4
|
|
>>> EPOCH_NUM = 4
|
|
|
|
>>> IMAGE_SIZE = 784
|
|
>>> CLASS_NUM = 10
|
|
|
|
>>> # define a random dataset
|
|
>>> class RandomDataset(paddle.io.Dataset): # type: ignore[type-arg]
|
|
... def __init__(self, num_samples):
|
|
... self.num_samples = num_samples
|
|
...
|
|
... def __getitem__(self, idx):
|
|
... image = np.random.random([IMAGE_SIZE]).astype('float32')
|
|
... label = np.random.randint(0, CLASS_NUM - 1, (1,)).astype('int64')
|
|
... return image, label
|
|
...
|
|
... def __len__(self):
|
|
... return self.num_samples
|
|
|
|
>>> paddle.enable_static()
|
|
|
|
>>> image = static.data(name='image', shape=[None, 784], dtype='float32')
|
|
>>> label = static.data(name='label', shape=[None, 1], dtype='int64')
|
|
>>> pred = static.nn.fc(x=image, size=10, activation='softmax')
|
|
>>> loss = F.cross_entropy(input=pred, label=label)
|
|
>>> avg_loss = paddle.mean(loss)
|
|
|
|
>>> optimizer = paddle.optimizer.SGD(learning_rate=0.001)
|
|
>>> optimizer.minimize(avg_loss)
|
|
|
|
>>> place = paddle.CPUPlace()
|
|
>>> exe = static.Executor(place)
|
|
>>> exe.run(static.default_startup_program())
|
|
|
|
>>> # create data loader
|
|
>>> dataset = RandomDataset(BATCH_NUM * BATCH_SIZE)
|
|
>>> loader = paddle.io.DataLoader(
|
|
... dataset,
|
|
... feed_list=[image, label],
|
|
... places=place,
|
|
... batch_size=BATCH_SIZE,
|
|
... shuffle=True,
|
|
... drop_last=True,
|
|
... return_list=False,
|
|
... num_workers=2,
|
|
... )
|
|
|
|
>>> # 1. train and save inference model
|
|
>>> for data in loader():
|
|
... exe.run(
|
|
... static.default_main_program(),
|
|
... feed=data,
|
|
... fetch_list=[avg_loss],
|
|
... )
|
|
|
|
>>> model_path = "fc.example.model"
|
|
>>> paddle.static.save_inference_model(
|
|
... model_path,
|
|
... [image],
|
|
... [pred],
|
|
... exe,
|
|
... )
|
|
|
|
>>> # 2. load model
|
|
|
|
>>> # enable dygraph mode
|
|
>>> paddle.disable_static(place)
|
|
|
|
>>> # load
|
|
>>> fc = paddle.jit.load(model_path)
|
|
|
|
>>> # inference
|
|
>>> fc.eval()
|
|
>>> x = paddle.randn([1, IMAGE_SIZE], 'float32')
|
|
>>> pred = fc(x)
|
|
|
|
>>> # fine-tune
|
|
>>> fc.train()
|
|
>>> loss_fn = nn.CrossEntropyLoss()
|
|
>>> adam = opt.Adam(learning_rate=0.001, parameters=fc.parameters())
|
|
>>> loader = paddle.io.DataLoader(
|
|
... dataset,
|
|
... places=place,
|
|
... batch_size=BATCH_SIZE,
|
|
... shuffle=True,
|
|
... drop_last=True,
|
|
... num_workers=2,
|
|
... )
|
|
>>> for epoch_id in range(EPOCH_NUM):
|
|
... for batch_id, (image, label) in enumerate(loader()):
|
|
... out = fc(image)
|
|
... loss = loss_fn(out, label)
|
|
... loss.backward()
|
|
... adam.step()
|
|
... adam.clear_grad()
|
|
... print("Epoch {} batch {}: loss = {}".format(epoch_id, batch_id, np.mean(loss.numpy())))
|
|
"""
|
|
# 1. construct correct config
|
|
config = _parse_load_config(configs)
|
|
model_path, config = _build_load_path_and_config(path, config)
|
|
|
|
if use_pir_api():
|
|
return PirTranslatedLayer._construct(model_path, config)
|
|
else:
|
|
return TranslatedLayer._construct(model_path, config)
|
|
|
|
|
|
def set_dynamic_shape(variable, shape_list):
|
|
if paddle.base.dygraph.base.in_to_static_mode():
|
|
if isinstance(variable, paddle.base.framework.Variable):
|
|
variable.desc.set_shape(shape_list)
|
|
elif isinstance(variable, paddle.pir.Value):
|
|
variable.set_shape(shape_list)
|
|
else:
|
|
raise TypeError(
|
|
"In to_static mode, variable must be a Variable or Value"
|
|
)
|
|
else:
|
|
# in dygraph mode, dynamic shape is not needed, just do nothing.
|
|
return
|
|
|
|
|
|
def get_ast_static_function(function):
|
|
if isinstance(function, SymbolicStaticFunction):
|
|
if function.class_instance:
|
|
dygraph_function = types.MethodType(
|
|
function._dygraph_function, function.class_instance
|
|
)
|
|
else:
|
|
dygraph_function = function._dygraph_function
|
|
|
|
if function._function_spec._input_spec is None:
|
|
ast_static_function = ASTStaticFunction(
|
|
dygraph_function,
|
|
function.last_call_input_spec,
|
|
**function._kwargs,
|
|
)
|
|
return ast_static_function
|
|
else:
|
|
ast_static_function = ASTStaticFunction(
|
|
dygraph_function,
|
|
function._function_spec._input_spec,
|
|
**function._kwargs,
|
|
)
|
|
return ast_static_function
|
|
return function
|
|
|
|
|
|
def json_to_pdmodel(net, input_spec, load_path, save_path):
|
|
net1 = paddle.jit.load(load_path)
|
|
state_dict = {}
|
|
for val in net1.state_dict().values():
|
|
name = val.name[: val.name.rfind('_')]
|
|
state_dict[name] = val
|
|
|
|
name_state_dict = {}
|
|
for name, var in net.to_static_state_dict().items():
|
|
name_state_dict[name] = state_dict[var.name]
|
|
|
|
net.set_state_dict(name_state_dict)
|
|
with paddle.pir_utils.OldIrGuard():
|
|
paddle.jit.save(net, save_path, input_spec)
|