Files
paddlepaddle--paddle/test/dygraph_to_static/transformer_util.py
T
2026-07-13 12:40:42 +08:00

436 lines
13 KiB
Python

# 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.
import pickle
import warnings
from functools import partial
import numpy as np
import paddle
from paddle.dataset import wmt16
def get_input_descs(args, mode="train"):
batch_size = args.batch_size # TODO None(before)
seq_len = None
n_head = getattr(args, "n_head", 8)
d_model = getattr(args, "d_model", 512)
input_descs_train = {
"src_word": [(batch_size, seq_len), "int64", 2],
"src_pos": [(batch_size, seq_len), "int64"],
"src_slf_attn_bias": [
(batch_size, n_head, seq_len, seq_len),
"float32",
],
"trg_word": [(batch_size, seq_len), "int64", 2],
"trg_pos": [(batch_size, seq_len), "int64"],
"trg_slf_attn_bias": [
(batch_size, n_head, seq_len, seq_len),
"float32",
],
"trg_src_attn_bias": [
(batch_size, n_head, seq_len, seq_len),
"float32",
], # TODO: 1 for predict, seq_len for train
"enc_output": [(batch_size, seq_len, d_model), "float32"],
"lbl_word": [(None, 1), "int64"],
"lbl_weight": [(None, 1), "float32"],
"init_score": [(batch_size, 1), "float32", 2],
"init_idx": [(batch_size,), "int32"],
}
input_descs_predict = {
"src_word": [(batch_size, seq_len), "int64", 2],
"src_pos": [(batch_size, seq_len), "int64"],
"src_slf_attn_bias": [
(batch_size, n_head, seq_len, seq_len),
"float32",
],
"trg_word": [(batch_size, seq_len), "int64", 2],
"trg_pos": [(batch_size, seq_len), "int64"],
"trg_slf_attn_bias": [
(batch_size, n_head, seq_len, seq_len),
"float32",
],
"trg_src_attn_bias": [(batch_size, n_head, 1, seq_len), "float32"],
"enc_output": [(batch_size, seq_len, d_model), "float32"],
"lbl_word": [(None, 1), "int64"],
"lbl_weight": [(None, 1), "float32"],
"init_score": [(batch_size, 1), "float32", 2],
"init_idx": [(batch_size,), "int32"],
}
return input_descs_train if mode == "train" else input_descs_predict
encoder_data_input_fields = (
"src_word",
"src_pos",
"src_slf_attn_bias",
)
decoder_data_input_fields = (
"trg_word",
"trg_pos",
"trg_slf_attn_bias",
"trg_src_attn_bias",
"enc_output",
)
label_data_input_fields = (
"lbl_word",
"lbl_weight",
)
fast_decoder_data_input_fields = (
"trg_word",
"trg_src_attn_bias",
)
class ModelHyperParams:
print_step = 2
save_dygraph_model_path = "dygraph_trained_models"
save_static_model_path = "static_trained_models"
inference_model_dir = "infer_model"
output_file = "predict.txt"
batch_size = 5
epoch = 1
learning_rate = 2.0
beta1 = 0.9
beta2 = 0.997
eps = 1e-9
warmup_steps = 8000
label_smooth_eps = 0.1
beam_size = 5
max_out_len = 5 # small number to avoid the unittest timeout
n_best = 1
src_vocab_size = 36556
trg_vocab_size = 36556
bos_idx = 0 # index for <bos> token
eos_idx = 1 # index for <eos> token
unk_idx = 2 # index for <unk> token
max_length = 256
d_model = 512
d_inner_hid = 2048
d_key = 64
d_value = 64
n_head = 8
n_layer = 6
prepostprocess_dropout = 0.1
attention_dropout = 0.1
relu_dropout = 0.1
preprocess_cmd = "n" # layer normalization
postprocess_cmd = "da" # dropout + residual connection
weight_sharing = True
def pad_batch_data(
insts,
pad_idx,
n_head,
is_target=False,
is_label=False,
return_attn_bias=True,
return_max_len=True,
return_num_token=False,
):
return_list = []
max_len = max(len(inst) for inst in insts)
inst_data = np.array(
[inst + [pad_idx] * (max_len - len(inst)) for inst in insts]
)
return_list += [inst_data.astype("int64").reshape([-1, 1])]
if is_label: # label weight
inst_weight = np.array(
[
[1.0] * len(inst) + [0.0] * (max_len - len(inst))
for inst in insts
]
)
return_list += [inst_weight.astype("float32").reshape([-1, 1])]
else: # position data
inst_pos = np.array(
[
list(range(0, len(inst))) + [0] * (max_len - len(inst))
for inst in insts
]
)
return_list += [inst_pos.astype("int64").reshape([-1, 1])]
if return_attn_bias:
if is_target:
slf_attn_bias_data = np.ones((inst_data.shape[0], max_len, max_len))
slf_attn_bias_data = np.triu(slf_attn_bias_data, 1).reshape(
[-1, 1, max_len, max_len]
)
slf_attn_bias_data = np.tile(
slf_attn_bias_data, [1, n_head, 1, 1]
) * [-1e9]
else:
slf_attn_bias_data = np.array(
[
[0] * len(inst) + [-1e9] * (max_len - len(inst))
for inst in insts
]
)
slf_attn_bias_data = np.tile(
slf_attn_bias_data.reshape([-1, 1, 1, max_len]),
[1, n_head, max_len, 1],
)
return_list += [slf_attn_bias_data.astype("float32")]
if return_max_len:
return_list += [max_len]
if return_num_token:
num_token = 0
for inst in insts:
num_token += len(inst)
return_list += [num_token]
return return_list if len(return_list) > 1 else return_list[0]
def prepare_train_input(insts, src_pad_idx, trg_pad_idx, n_head):
src_word, src_pos, src_slf_attn_bias, src_max_len = pad_batch_data(
[inst[0] for inst in insts], src_pad_idx, n_head, is_target=False
)
src_word = src_word.reshape(-1, src_max_len)
src_pos = src_pos.reshape(-1, src_max_len)
trg_word, trg_pos, trg_slf_attn_bias, trg_max_len = pad_batch_data(
[inst[1] for inst in insts], trg_pad_idx, n_head, is_target=True
)
trg_word = trg_word.reshape(-1, trg_max_len)
trg_pos = trg_pos.reshape(-1, trg_max_len)
trg_src_attn_bias = np.tile(
src_slf_attn_bias[:, :, ::src_max_len, :], [1, 1, trg_max_len, 1]
).astype("float32")
lbl_word, lbl_weight, num_token = pad_batch_data(
[inst[2] for inst in insts],
trg_pad_idx,
n_head,
is_target=False,
is_label=True,
return_attn_bias=False,
return_max_len=False,
return_num_token=True,
)
lbl_word = lbl_word.reshape(-1, 1)
lbl_weight = lbl_weight.reshape(-1, 1)
data_inputs = [
src_word,
src_pos,
src_slf_attn_bias,
trg_word,
trg_pos,
trg_slf_attn_bias,
trg_src_attn_bias,
lbl_word,
lbl_weight,
]
return data_inputs
def prepare_infer_input(insts, src_pad_idx, bos_idx, n_head):
src_word, src_pos, src_slf_attn_bias, src_max_len = pad_batch_data(
[inst[0] for inst in insts], src_pad_idx, n_head, is_target=False
)
# start tokens
trg_word = np.asarray([[bos_idx]] * len(insts), dtype="int64")
trg_src_attn_bias = np.tile(
src_slf_attn_bias[:, :, ::src_max_len, :], [1, 1, 1, 1]
).astype("float32")
trg_word = trg_word.reshape(-1, 1)
src_word = src_word.reshape(-1, src_max_len)
src_pos = src_pos.reshape(-1, src_max_len)
data_inputs = [
src_word,
src_pos,
src_slf_attn_bias,
trg_word,
trg_src_attn_bias,
]
return data_inputs
def get_feed_data_reader(args, mode='train'):
def __for_train__():
train_reader = paddle.batch(
wmt16.train(args.src_vocab_size, args.trg_vocab_size),
batch_size=args.batch_size,
)
for batch in train_reader():
tensors = prepare_train_input(
batch, args.eos_idx, args.eos_idx, args.n_head
)
yield tensors
def __for_test__():
test_reader = paddle.batch(
wmt16.test(args.src_vocab_size, args.trg_vocab_size),
batch_size=args.batch_size,
)
for batch in test_reader():
tensors = prepare_infer_input(
batch, args.eos_idx, args.eos_idx, args.n_head
)
yield tensors
return __for_train__ if mode == 'train' else __for_test__
class InputField:
def __init__(self, input_slots):
self.feed_list = []
for slot in input_slots:
self.feed_list.append(
paddle.static.data(
name=slot['name'],
shape=slot['shape'],
dtype=slot['dtype'],
lod_level=slot.get('lod_level', 0),
)
)
class TransedWMT16TrainDataSet(paddle.io.Dataset):
def __init__(self, data_reader, length):
self.src_word = []
self.src_pos = []
self.src_slf_attn_bias = []
self.trg_word = []
self.trg_pos = []
self.trg_slf_attn_bias = []
self.trg_src_attn_bias = []
self.lbl_word = []
self.lbl_weight = []
self.reader = data_reader()
self._generate(length)
def _generate(self, length):
for i, data in enumerate(self.reader):
if i >= length:
break
self.src_word.append(data[0])
self.src_pos.append(data[1])
self.src_slf_attn_bias.append(data[2])
self.trg_word.append(data[3])
self.trg_pos.append(data[4])
self.trg_slf_attn_bias.append(data[5])
self.trg_src_attn_bias.append(data[6])
self.lbl_word.append(data[7])
self.lbl_weight.append(data[8])
def __getitem__(self, idx):
return (
self.src_word[idx],
self.src_pos[idx],
self.src_slf_attn_bias[idx],
self.trg_word[idx],
self.trg_pos[idx],
self.trg_slf_attn_bias[idx],
self.trg_src_attn_bias[idx],
self.lbl_word[idx],
self.lbl_weight[idx],
)
def __len__(self):
return len(self.src_word)
class TransedWMT16TestDataSet(paddle.io.Dataset):
def __init__(self, data_reader, length):
self.src_word = []
self.src_pos = []
self.src_slf_attn_bias = []
self.trg_word = []
self.trg_pos = []
self.trg_slf_attn_bias = []
self.trg_src_attn_bias = []
self.lbl_word = []
self.lbl_weight = []
self.reader = data_reader()
self._generate(length)
def _generate(self, length):
for i, data in enumerate(self.reader):
if i >= length:
break
self.src_word.append(data[0])
self.src_pos.append(data[1])
self.src_slf_attn_bias.append(data[2])
self.trg_word.append(data[3])
self.trg_slf_attn_bias.append(data[4])
def __getitem__(self, idx):
return (
self.src_word[idx],
self.src_pos[idx],
self.src_slf_attn_bias[idx],
self.trg_word[idx],
self.trg_slf_attn_bias[idx],
)
def __len__(self):
return len(self.src_word)
def load(program, model_path, executor=None, var_list=None):
"""
To load python2 saved models in python3.
"""
try:
paddle.static.load(program, model_path, executor, var_list)
except UnicodeDecodeError:
warnings.warn(
"An UnicodeDecodeError is caught, which might be caused by loading "
"a python2 saved model. Encoding of pickle.load would be set and "
"load again automatically."
)
load_bak = pickle.load
pickle.load = partial(load_bak, encoding="latin1")
paddle.static.load(program, model_path, executor, var_list)
pickle.load = load_bak
def load_dygraph(model_path, keep_name_table=False):
"""
To load python2 saved models in python3.
"""
try:
para_dict = paddle.load(
model_path + '.pdparams', keep_name_table=keep_name_table
)
opti_dict = paddle.load(
model_path + '.pdopt', keep_name_table=keep_name_table
)
return para_dict, opti_dict
except UnicodeDecodeError:
warnings.warn(
"An UnicodeDecodeError is caught, which might be caused by loading "
"a python2 saved model. Encoding of pickle.load would be set and "
"load again automatically."
)
load_bak = pickle.load
pickle.load = partial(load_bak, encoding="latin1")
para_dict = paddle.load(
model_path + '.pdparams', keep_name_table=keep_name_table
)
opti_dict = paddle.load(
model_path + '.pdopt', keep_name_table=keep_name_table
)
pickle.load = load_bak
return para_dict, opti_dict