chore: import upstream snapshot with attribution
This commit is contained in:
@@ -0,0 +1,797 @@
|
||||
# Copyright (c) 2020 PaddlePaddle Authors. 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 os
|
||||
|
||||
import numpy as np
|
||||
|
||||
import paddle
|
||||
from paddle.base import core, framework, unique_name
|
||||
from paddle.base.dygraph.base import switch_to_static_graph
|
||||
from paddle.framework import in_dynamic_mode
|
||||
from paddle.nn.layer import layers
|
||||
from paddle.pir.core import datatype_to_vartype
|
||||
|
||||
__all__ = []
|
||||
|
||||
PIR_INFER_MODEL_SUFFIX = ".json"
|
||||
|
||||
from .translated_layer import (
|
||||
BUFFER_NAME_PREFIX,
|
||||
INFER_PARAMS_SUFFIX,
|
||||
PARAMETER_NAME_PREFIX,
|
||||
)
|
||||
|
||||
|
||||
def _load_pir_program(model_file_path):
|
||||
program = paddle.static.Program()
|
||||
trainable = paddle.base.core.deserialize_pir_program(
|
||||
model_file_path, program
|
||||
)
|
||||
|
||||
return program, trainable
|
||||
|
||||
|
||||
@switch_to_static_graph
|
||||
def _generate_unique_var_name(prefix):
|
||||
return unique_name.generate_with_ignorable_key(prefix)
|
||||
|
||||
|
||||
@switch_to_static_graph
|
||||
def _generate_unique_var_name(prefix):
|
||||
return unique_name.generate(prefix)
|
||||
|
||||
|
||||
from paddle.static.pir_io import get_pir_parameters
|
||||
|
||||
|
||||
def _get_pir_parameters_var_names(program):
|
||||
persistable_vars = []
|
||||
persistable_names = []
|
||||
rename_new_old_dict = {}
|
||||
param, opt = get_pir_parameters(program)
|
||||
vars = param + opt
|
||||
for var in vars:
|
||||
persistable_vars.append(var)
|
||||
origin_name = var.name
|
||||
var.name = _generate_unique_var_name(var.name)
|
||||
rename_new_old_dict[var.name] = origin_name
|
||||
persistable_names.append(var.name)
|
||||
|
||||
return (
|
||||
persistable_vars,
|
||||
persistable_names,
|
||||
rename_new_old_dict,
|
||||
)
|
||||
|
||||
|
||||
class _PirProgramHolder:
|
||||
def __init__(self, program, trainable):
|
||||
super().__init__()
|
||||
|
||||
# input, output, persistable,
|
||||
self._input_vars = []
|
||||
self._output_vars = []
|
||||
self._parameter_vars = []
|
||||
self._parameter_names = []
|
||||
|
||||
self.support_train = trainable
|
||||
# append suffix var name dict
|
||||
self._suffix_varname_dict = None
|
||||
self._infer_program = program
|
||||
|
||||
self._preprocess()
|
||||
|
||||
def _preprocess(self):
|
||||
(
|
||||
self._parameter_vars,
|
||||
self._parameter_names,
|
||||
self._suffix_varname_dict,
|
||||
) = _get_pir_parameters_var_names(self._infer_program)
|
||||
block = self._infer_program.global_block()
|
||||
for op in block.ops:
|
||||
if op.name() == 'pd_op.data':
|
||||
self._input_vars.append(op.result(0))
|
||||
elif op.name() == 'pd_op.feed':
|
||||
var_name = op.attr()["name"]
|
||||
org_value = op.result(0)
|
||||
with block:
|
||||
value = paddle._pir_ops.data(
|
||||
name=var_name,
|
||||
shape=org_value.shape,
|
||||
dtype=org_value.dtype,
|
||||
place=paddle.base.core.Place(),
|
||||
)
|
||||
org_value.replace_all_uses_with(value)
|
||||
value.get_defining_op().move_before(op)
|
||||
block.remove_op(op)
|
||||
if op.name() == 'pd_op.fetch':
|
||||
self._output_vars.append(op.operand_source(0))
|
||||
with block:
|
||||
paddle._pir_ops.set_persistable_value(
|
||||
op.operand_source(0),
|
||||
"output_" + str(len(self._output_vars) - 1),
|
||||
)
|
||||
block.remove_op(op)
|
||||
|
||||
@property
|
||||
def infer_program(self):
|
||||
return self._infer_program
|
||||
|
||||
@property
|
||||
def input_vars(self):
|
||||
return self._input_vars
|
||||
|
||||
@property
|
||||
def output_vars(self):
|
||||
return self._output_vars
|
||||
|
||||
@property
|
||||
def persistable_names(self):
|
||||
return self._parameter_names
|
||||
|
||||
@property
|
||||
def persistable_vars(self):
|
||||
return self._parameter_vars
|
||||
|
||||
|
||||
# [ PirTranslatedLayer : Run program in dygraph mode ]
|
||||
#
|
||||
# DESIGN IDEA: using an special operator `PirRunProgram`, execute program inside operator.
|
||||
#
|
||||
# Op's Inputs:
|
||||
# - the input variable of the user feed
|
||||
# - the necessary parameters of the network
|
||||
# Op's Outputs:
|
||||
# - the output variable of fetch
|
||||
#
|
||||
# This op receives a complete program desc, internally creates scope
|
||||
# and executor, executes this program. Key points:
|
||||
#
|
||||
# 1. Data Sharing:
|
||||
# The variable/parameter of the dynamic graph is not in the scope, so before the op
|
||||
# executes the program internally, create persistent variables with the
|
||||
# same name as feed, parameters, and fetch in the scope, and share the
|
||||
# DenseTensor of the op input.
|
||||
#
|
||||
# 2. Forward and Backward Separation:
|
||||
# Because the dynamic graph op performs the forward and backward separately,
|
||||
# in the forward op RunProgram, we only execute the forward part of whole program,
|
||||
# and in the backward op RunProgramGrad, we execute the backward part of program.
|
||||
# We can not separate the program into forward and backward part, which will
|
||||
# make some control flow execution logic wrong.
|
||||
|
||||
|
||||
# NOTE: [compatible] deal with model saved by save_inference_model,
|
||||
# which need get var info from program desc
|
||||
|
||||
|
||||
def _load_pir_parameter_vars(model_path, program_holder, params_filename):
|
||||
# construct var dict
|
||||
load_var_dict = {}
|
||||
load_var_list = []
|
||||
other_var_dict = {}
|
||||
load_densetensor_list = []
|
||||
persistable_var = program_holder.persistable_vars
|
||||
persistable_var_name = program_holder.persistable_names
|
||||
origin_persistable_var_name = [
|
||||
program_holder._suffix_varname_dict[var_name]
|
||||
for var_name in persistable_var_name
|
||||
]
|
||||
for name, var in sorted(zip(origin_persistable_var_name, persistable_var)):
|
||||
if var.persistable:
|
||||
# use default shape and dtype
|
||||
new_var = framework.EagerParamBase(
|
||||
shape=var.shape, # only to pass check, this shape is not meaningful
|
||||
dtype=core.VarDesc.VarType.FP32,
|
||||
name=var.name,
|
||||
persistable=True,
|
||||
)
|
||||
|
||||
new_var.stop_gradient = var.stop_gradient
|
||||
load_var_dict[name] = new_var
|
||||
load_var_list.append(new_var)
|
||||
load_densetensor_list.append(new_var.get_tensor())
|
||||
|
||||
else:
|
||||
new_var = core.eager.Tensor(
|
||||
dtype=datatype_to_vartype[var.dtype],
|
||||
dims=var.shape,
|
||||
name=var.name,
|
||||
type=core.VarDesc.VarType.DENSE_TENSOR,
|
||||
place=framework._current_expected_place(),
|
||||
persistable=False,
|
||||
)
|
||||
other_var_dict[name] = new_var
|
||||
|
||||
# load all vars
|
||||
assert params_filename is not None, "params_filename should not be None."
|
||||
var_file_path = os.path.join(model_path, params_filename)
|
||||
if os.path.exists(var_file_path):
|
||||
core.load_combine_func(
|
||||
var_file_path,
|
||||
list(load_var_dict.keys()),
|
||||
load_densetensor_list,
|
||||
False,
|
||||
framework._current_expected_place(),
|
||||
)
|
||||
else:
|
||||
raise ValueError(
|
||||
f"The file {var_file_path} does not exist. Please check the model path."
|
||||
)
|
||||
|
||||
load_var_dict.update(other_var_dict)
|
||||
return load_var_dict
|
||||
|
||||
|
||||
def _construct_program_holders(model_path, model_filename=None):
|
||||
# make sure the path has been checked
|
||||
program_holder_dict = {}
|
||||
|
||||
if model_filename is not None:
|
||||
# [compatible] if assign model_filename, only can load one program as Layer.forward
|
||||
model_filename = os.path.basename(model_filename)
|
||||
model_file_path = os.path.join(model_path, model_filename)
|
||||
model_name = model_filename[: -len(PIR_INFER_MODEL_SUFFIX)]
|
||||
# Load every file that meets the requirements in the directory model_path.
|
||||
for filename in os.listdir(model_path):
|
||||
if model_filename == filename:
|
||||
func_name = 'forward'
|
||||
model_file_path = os.path.join(model_path, model_filename)
|
||||
elif filename.endswith(
|
||||
PIR_INFER_MODEL_SUFFIX
|
||||
) and filename.startswith(model_name):
|
||||
parsing_names = filename[
|
||||
len(model_name) : -len(PIR_INFER_MODEL_SUFFIX) + 1
|
||||
].split('.')
|
||||
if len(parsing_names) == 3 and len(parsing_names[1]) > 0:
|
||||
func_name = parsing_names[1]
|
||||
model_file_path = os.path.join(model_path, filename)
|
||||
else:
|
||||
continue
|
||||
else:
|
||||
continue
|
||||
program, trainable = _load_pir_program(model_file_path)
|
||||
program_holder_dict[func_name] = _PirProgramHolder(
|
||||
program, trainable
|
||||
)
|
||||
else:
|
||||
for _, _, file_names in os.walk(model_path):
|
||||
for name in file_names:
|
||||
if 'model' in name:
|
||||
model_file_path = os.path.join(model_path, name)
|
||||
method_name = name.strip('_')
|
||||
if method_name == 'model':
|
||||
method_name = 'forward'
|
||||
else:
|
||||
method_name.replace('model', '')
|
||||
program, trainable = _load_pir_program(model_file_path)
|
||||
program_holder_dict[method_name] = _PirProgramHolder(
|
||||
program, trainable
|
||||
)
|
||||
|
||||
return program_holder_dict
|
||||
|
||||
|
||||
def _construct_params_and_buffers(model_path, programs, params_filename=None):
|
||||
params_path = os.path.join(model_path, str(params_filename))
|
||||
|
||||
if params_filename is not None and not os.path.exists(params_path):
|
||||
# When saving XX, there is only '*.pdmodel'
|
||||
return {}
|
||||
else:
|
||||
var_dict = _load_pir_parameter_vars(
|
||||
model_path, programs['forward'], params_filename
|
||||
)
|
||||
model_name = params_filename[: -len(INFER_PARAMS_SUFFIX)]
|
||||
# Load every file that meets the requirements in the directory model_path.
|
||||
for file_name in os.listdir(model_path):
|
||||
if file_name.startswith(model_name) and file_name.endswith(
|
||||
INFER_PARAMS_SUFFIX
|
||||
):
|
||||
parsing_names = file_name[
|
||||
len(model_name) : -len(INFER_PARAMS_SUFFIX) + 1
|
||||
].split('.')
|
||||
if len(parsing_names) == 3 and len(parsing_names[1]) > 0:
|
||||
func_name = parsing_names[1]
|
||||
else:
|
||||
continue
|
||||
else:
|
||||
continue
|
||||
|
||||
var_dict.update(
|
||||
_load_pir_parameter_vars(
|
||||
model_path, programs[func_name], file_name
|
||||
)
|
||||
)
|
||||
|
||||
return var_dict
|
||||
|
||||
|
||||
def _run_dygraph(instance, input, program_holder, method_name):
|
||||
# 1. prepare inputs, outputs, attrs
|
||||
input_tensors = []
|
||||
input_tensor_names = []
|
||||
|
||||
for i, value in enumerate(input):
|
||||
if not isinstance(value, (np.ndarray, core.eager.Tensor)):
|
||||
raise TypeError(
|
||||
f"The type of input in PirTranslatedLayer must be numpy array or Variable(Tensor), but received {type(value)}."
|
||||
)
|
||||
# NOTE: In order to unify the API, firstly convert the input to Tensor
|
||||
if isinstance(value, np.ndarray):
|
||||
tensor = core.eager.Tensor(
|
||||
value=value,
|
||||
name=program_holder.input_vars[i].name,
|
||||
persistable=False,
|
||||
place=framework._current_expected_place(),
|
||||
zero_copy=True,
|
||||
)
|
||||
else:
|
||||
tensor = value
|
||||
# NOTE: we changed var name here,
|
||||
# but it may be an important name set by user
|
||||
tensor.name = program_holder.input_vars[i].name
|
||||
input_tensor_names.append(tensor.name)
|
||||
input_tensors.append(tensor)
|
||||
|
||||
if instance._get_partial_program_layer(method_name) is None:
|
||||
persistable_tensors = []
|
||||
origin_persistable_var_name = [
|
||||
program_holder._suffix_varname_dict[var_name]
|
||||
for var_name in program_holder.persistable_names
|
||||
]
|
||||
for var_name in origin_persistable_var_name:
|
||||
dy_var_name = instance._persistable_var_name_dict[var_name]
|
||||
if dy_var_name in instance._parameters:
|
||||
persistable_tensors.append(instance._parameters[dy_var_name])
|
||||
elif dy_var_name in instance._buffers:
|
||||
persistable_tensors.append(instance._buffers[dy_var_name])
|
||||
else:
|
||||
raise ValueError(
|
||||
f"The persistable variable {var_name} does not exist in current PirTranslatedLayer."
|
||||
)
|
||||
|
||||
from paddle.jit.dy2static.pir_partial_program import PartialProgramLayer
|
||||
|
||||
inputs = program_holder.input_vars
|
||||
outputs = program_holder.output_vars
|
||||
parameters = (persistable_tensors, program_holder.persistable_vars)
|
||||
|
||||
layer = PartialProgramLayer(
|
||||
program_holder.infer_program,
|
||||
inputs,
|
||||
outputs,
|
||||
parameters,
|
||||
)
|
||||
instance._set_partial_program_layer(method_name, layer)
|
||||
layer = instance._get_partial_program_layer(method_name)
|
||||
if instance._is_test:
|
||||
layer.training = False
|
||||
else:
|
||||
if not program_holder.support_train:
|
||||
raise ValueError(
|
||||
"The model is not trainable, please check model_file of jit.save."
|
||||
)
|
||||
else:
|
||||
layer.training = True
|
||||
|
||||
return layer(input_tensors)
|
||||
|
||||
|
||||
def _run_static_graph(inputs, program_holder, src_program):
|
||||
'''
|
||||
This function is used when the pirTranslatedLayer is
|
||||
applied for dy_to_static conversion.
|
||||
'''
|
||||
dst_program = paddle.static.default_main_program()
|
||||
value_map = paddle.pir.IrMapping()
|
||||
# Establish a mapping relationship between existing parameters
|
||||
# and corresponding parameters in the program to be copied
|
||||
len_dst_op = len(dst_program.global_block().ops)
|
||||
for dst_op in dst_program.global_block().ops:
|
||||
if dst_op.name() == "builtin.parameter":
|
||||
for src_op in src_program.global_block().ops[:len_dst_op]:
|
||||
if (
|
||||
src_op.name() == dst_op.name()
|
||||
and src_op.result(0).name == dst_op.result(0).name
|
||||
):
|
||||
for i in range(src_op.num_results()):
|
||||
value_map.add(src_op.result(i), dst_op.result(i))
|
||||
# Establish a mapping relationship between truly inputs
|
||||
# and corresponding inputs in the program to be copied
|
||||
src_inputs = program_holder.input_vars
|
||||
if len(src_inputs) != len(inputs):
|
||||
raise ValueError(
|
||||
f"The number of input is invalid, expected {len(src_inputs)}, but received {len(inputs)}."
|
||||
)
|
||||
for src_input, input_ in zip(src_inputs, inputs):
|
||||
value_map.add(src_input, input_)
|
||||
|
||||
# find the insert point for copy
|
||||
current_insert_point = paddle.pir.get_current_insertion_point()
|
||||
current_block = current_insert_point.block()
|
||||
src_program.copy_to_block(value_map, current_block)
|
||||
|
||||
output = [value_map.look_up(v) for v in program_holder.output_vars]
|
||||
return output[0] if len(output) == 1 else output
|
||||
|
||||
|
||||
def _collect_current_and_parent_var(program, block_idx):
|
||||
'''
|
||||
Get variables in current block and its parent block.
|
||||
|
||||
Args:
|
||||
program(Program): The program containing the current block.
|
||||
block_idx(int): index of current block.
|
||||
|
||||
Returns:
|
||||
List: list of variables.
|
||||
'''
|
||||
vars = []
|
||||
if block_idx < 0:
|
||||
return vars
|
||||
for var in program.block(block_idx).vars:
|
||||
vars.append(var)
|
||||
parent_idx = program.block(block_idx).parent_idx
|
||||
if parent_idx > -1:
|
||||
vars += _collect_current_and_parent_var(program, parent_idx)
|
||||
return vars
|
||||
|
||||
|
||||
class PirTranslatedLayer(layers.Layer):
|
||||
"""
|
||||
PirTranslatedLayer is a ``paddle.nn.Layer`` for holding the model
|
||||
loaded by :ref:`api_paddle_jit_load` . It can be used like a
|
||||
general Layer object in eval or train mode.
|
||||
|
||||
.. note:
|
||||
The PirTranslatedLayer objects should not be created by constructor, it only can be loaded and constructed by :ref:`api_paddle_jit_load` .
|
||||
|
||||
Examples:
|
||||
.. code-block:: pycon
|
||||
|
||||
>>> # 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
|
||||
>>> model_path = "linear.example.model"
|
||||
>>> paddle.jit.save(layer, model_path)
|
||||
|
||||
>>> # 2. load model as PirTranslatedLayer
|
||||
>>> # load
|
||||
>>> translated_layer = paddle.jit.load(model_path)
|
||||
|
||||
>>> # inference
|
||||
>>> translated_layer.eval()
|
||||
>>> x = paddle.randn([1, IMAGE_SIZE], 'float32')
|
||||
>>> pred = translated_layer(x)
|
||||
|
||||
>>> # fine-tune
|
||||
>>> translated_layer.train()
|
||||
>>> adam = opt.Adam(learning_rate=0.001, parameters=translated_layer.parameters())
|
||||
>>> train(translated_layer, loader, loss_fn, adam)
|
||||
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
programs: dict[str, paddle.static.Program],
|
||||
persistable_vars: dict[str, paddle.Tensor],
|
||||
):
|
||||
super().__init__()
|
||||
|
||||
if not isinstance(programs, dict):
|
||||
raise TypeError(
|
||||
"PirTranslatedLayer need to use _ProgramHolder's dict for initialization."
|
||||
)
|
||||
if not isinstance(persistable_vars, dict):
|
||||
raise TypeError(
|
||||
"PirTranslatedLayer need to use persistable variable dict for initialization."
|
||||
)
|
||||
|
||||
self._program_holder_dict = programs
|
||||
|
||||
# NOTE(chenweihang): [ why not use var name directly? ]
|
||||
# When add parameter or buffer to Layer by follow apis,
|
||||
# the variable name can't contain `.`, because which may cause
|
||||
# AttributeError when access the newly added parameter or buffer
|
||||
# in the form of `self.**.**``, but the EagerParamBase or BarBase
|
||||
# name contains `.` originally, such as `linear_0.w_0`, so here
|
||||
# need to generate new var name for each var
|
||||
self._persistable_var_name_dict = {}
|
||||
# the PirTranslatedLayer object held var names count started from 0
|
||||
with unique_name.guard():
|
||||
for name, var in persistable_vars.items():
|
||||
if isinstance(var, framework.EagerParamBase):
|
||||
dy_name = _generate_unique_var_name(PARAMETER_NAME_PREFIX)
|
||||
self._persistable_var_name_dict[name] = dy_name
|
||||
self.add_parameter(dy_name, var)
|
||||
elif isinstance(var, core.eager.Tensor):
|
||||
dy_name = _generate_unique_var_name(BUFFER_NAME_PREFIX)
|
||||
self._persistable_var_name_dict[name] = dy_name
|
||||
self.register_buffer(dy_name, var)
|
||||
else:
|
||||
raise TypeError(
|
||||
"Adding persistent variable which to layer is not supported now"
|
||||
)
|
||||
|
||||
self._is_test = True
|
||||
self._input_args_names = None
|
||||
self._partial_program_layers = {}
|
||||
|
||||
@staticmethod
|
||||
@framework.dygraph_only
|
||||
def _construct(model_path, configs=None):
|
||||
# 0. dir and filename check
|
||||
model_path = os.path.normpath(model_path)
|
||||
if not os.path.isdir(model_path):
|
||||
raise ValueError(f"There is no directory named '{model_path}'")
|
||||
model_filename = None
|
||||
params_filename = None
|
||||
if configs is not None:
|
||||
model_filename = configs.model_filename
|
||||
params_filename = configs.params_filename
|
||||
|
||||
# 1. load program desc & construct _ProgramHolder
|
||||
programs = _construct_program_holders(model_path, model_filename)
|
||||
|
||||
# 2. load layer parameters & buffers
|
||||
persistable_vars = _construct_params_and_buffers(
|
||||
model_path, programs, params_filename
|
||||
)
|
||||
|
||||
# 3. construct PirTranslatedLayer object
|
||||
translated_layer = PirTranslatedLayer(programs, persistable_vars)
|
||||
|
||||
# 4. create PirTranslatedLayer's execution method
|
||||
for method_name, program_holder in programs.items():
|
||||
if translated_layer._input_args_names is None:
|
||||
translated_layer._input_args_names = [
|
||||
ins.name for ins in program_holder.input_vars
|
||||
]
|
||||
setattr(
|
||||
PirTranslatedLayer,
|
||||
method_name,
|
||||
PirTranslatedLayer._execution_method_creator(
|
||||
method_name, program_holder
|
||||
),
|
||||
)
|
||||
|
||||
# 5. set PirTranslatedLayer's default mode to eval
|
||||
translated_layer.eval()
|
||||
|
||||
return translated_layer
|
||||
|
||||
@staticmethod
|
||||
def _execution_method_creator(method_name, program_holder):
|
||||
def __i_m_p_l__(self, *input):
|
||||
program_holder = self._program_holder_dict[__i_m_p_l__.__name__]
|
||||
# When using jit.save, it runs in static graph mode.
|
||||
# Run in dynamic graph mode when the model is inferring.
|
||||
if in_dynamic_mode():
|
||||
return _run_dygraph(self, input, program_holder, method_name)
|
||||
else:
|
||||
return _run_static_graph(
|
||||
input, program_holder, program_holder.infer_program
|
||||
)
|
||||
|
||||
__i_m_p_l__.__name__ = method_name
|
||||
return __i_m_p_l__
|
||||
|
||||
def train(self):
|
||||
self._is_test = False
|
||||
self.training = True
|
||||
|
||||
def eval(self):
|
||||
self._is_test = True
|
||||
self.training = False
|
||||
|
||||
def program(self, method_name='forward'):
|
||||
"""
|
||||
Gets translated program of specified method.
|
||||
|
||||
Args:
|
||||
- method_name (string): method name corresponding to the program
|
||||
to be obtained. Default: 'forward'.
|
||||
|
||||
Returns:
|
||||
Program
|
||||
|
||||
Examples:
|
||||
.. code-block:: pycon
|
||||
|
||||
>>> # doctest: +SKIP('`paddle.jit.to_static` can not run in xdoctest')
|
||||
>>> import numpy as np
|
||||
>>> import paddle
|
||||
>>> from paddle import 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())))
|
||||
>>> # 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
|
||||
>>> model_path = "linear.example.model"
|
||||
>>> paddle.jit.save(layer, model_path)
|
||||
|
||||
>>> # load
|
||||
>>> translated_layer = paddle.jit.load(model_path)
|
||||
|
||||
>>> # get program
|
||||
>>> program = translated_layer.program()
|
||||
"""
|
||||
# 1. get program holder
|
||||
program_holder = self._get_program_holder(method_name)
|
||||
|
||||
# 2. get inference program desc
|
||||
program = program_holder.infer_program
|
||||
|
||||
return program
|
||||
|
||||
def _get_program_holder(self, method_name='forward'):
|
||||
program_holder = self._program_holder_dict.get(method_name, None)
|
||||
if program_holder is None:
|
||||
raise ValueError(
|
||||
f"The method `{method_name}` does not exist in loaded PirTranslatedLayer."
|
||||
)
|
||||
return program_holder
|
||||
|
||||
def _input_spec(self, method_name='forward'):
|
||||
# 1. get program holder
|
||||
program_holder = self._get_program_holder(method_name)
|
||||
|
||||
# 2. build input spec by input desc
|
||||
input_spec = []
|
||||
for var in program_holder.input_vars:
|
||||
spec = paddle.static.InputSpec(
|
||||
shape=var.shape,
|
||||
dtype=var.dtype,
|
||||
name=var.name,
|
||||
)
|
||||
input_spec.append(spec)
|
||||
|
||||
return input_spec
|
||||
|
||||
def _output_spec(self, method_name='forward'):
|
||||
# 1. get program holder
|
||||
program_holder = self._get_program_holder(method_name)
|
||||
|
||||
# 2. build output spec by output desc
|
||||
output_spec = []
|
||||
for var in program_holder.output_vars:
|
||||
# NOTE(chenweihang): InputSpec describes a tensor, not just input.
|
||||
# Maybe the name is not good enough. Here we use InputSpec to
|
||||
# construct the description of Output tensor
|
||||
spec = paddle.static.InputSpec(
|
||||
shape=var.shape,
|
||||
dtype=var.dtype,
|
||||
name=var.name,
|
||||
)
|
||||
output_spec.append(spec)
|
||||
|
||||
return output_spec
|
||||
|
||||
def _get_partial_program_layer(self, method_name):
|
||||
return self._partial_program_layers.get(method_name, None)
|
||||
|
||||
def _set_partial_program_layer(self, method_name, layer):
|
||||
self._partial_program_layers[method_name] = layer
|
||||
Reference in New Issue
Block a user