325 lines
11 KiB
Python
325 lines
11 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 logging
|
|
import os
|
|
import tempfile
|
|
import time
|
|
import unittest
|
|
|
|
import numpy as np
|
|
import transformer_util as util
|
|
from dygraph_to_static_utils import (
|
|
Dy2StTestBase,
|
|
enable_to_static_guard,
|
|
test_default_mode_only,
|
|
)
|
|
from transformer_dygraph_model import (
|
|
CrossEntropyCriterion,
|
|
Transformer,
|
|
position_encoding_init,
|
|
)
|
|
|
|
import paddle
|
|
|
|
trainer_count = 1
|
|
place = (
|
|
paddle.CUDAPlace(0) if paddle.is_compiled_with_cuda() else paddle.CPUPlace()
|
|
)
|
|
SEED = 10
|
|
STEP_NUM = 10
|
|
|
|
|
|
def train_dygraph(args, batch_generator):
|
|
if SEED is not None:
|
|
paddle.seed(SEED)
|
|
paddle.framework.random._manual_program_seed(SEED)
|
|
# define data loader
|
|
|
|
train_loader = paddle.io.DataLoader(
|
|
batch_generator, batch_size=None, places=place
|
|
)
|
|
|
|
# define model
|
|
transformer = paddle.jit.to_static(
|
|
Transformer(
|
|
args.src_vocab_size,
|
|
args.trg_vocab_size,
|
|
args.max_length + 1,
|
|
args.n_layer,
|
|
args.n_head,
|
|
args.d_key,
|
|
args.d_value,
|
|
args.d_model,
|
|
args.d_inner_hid,
|
|
args.prepostprocess_dropout,
|
|
args.attention_dropout,
|
|
args.relu_dropout,
|
|
args.preprocess_cmd,
|
|
args.postprocess_cmd,
|
|
args.weight_sharing,
|
|
args.bos_idx,
|
|
args.eos_idx,
|
|
)
|
|
)
|
|
# define loss
|
|
criterion = CrossEntropyCriterion(args.label_smooth_eps)
|
|
# define optimizer
|
|
learning_rate = paddle.optimizer.lr.NoamDecay(
|
|
args.d_model, args.warmup_steps, args.learning_rate
|
|
)
|
|
# define optimizer
|
|
optimizer = paddle.optimizer.Adam(
|
|
learning_rate=learning_rate,
|
|
beta1=args.beta1,
|
|
beta2=args.beta2,
|
|
epsilon=float(args.eps),
|
|
parameters=transformer.parameters(),
|
|
)
|
|
# the best cross-entropy value with label smoothing
|
|
loss_normalizer = -(
|
|
(1.0 - args.label_smooth_eps) * np.log(1.0 - args.label_smooth_eps)
|
|
+ args.label_smooth_eps
|
|
* np.log(args.label_smooth_eps / (args.trg_vocab_size - 1) + 1e-20)
|
|
)
|
|
ce_time = []
|
|
ce_ppl = []
|
|
avg_loss = []
|
|
step_idx = 0
|
|
for pass_id in range(args.epoch):
|
|
pass_start_time = time.time()
|
|
batch_id = 0
|
|
for input_data in train_loader():
|
|
(
|
|
src_word,
|
|
src_pos,
|
|
src_slf_attn_bias,
|
|
trg_word,
|
|
trg_pos,
|
|
trg_slf_attn_bias,
|
|
trg_src_attn_bias,
|
|
lbl_word,
|
|
lbl_weight,
|
|
) = input_data
|
|
logits = transformer(
|
|
src_word,
|
|
src_pos,
|
|
src_slf_attn_bias,
|
|
trg_word,
|
|
trg_pos,
|
|
trg_slf_attn_bias,
|
|
trg_src_attn_bias,
|
|
)
|
|
sum_cost, avg_cost, token_num = criterion(
|
|
logits, lbl_word, lbl_weight
|
|
)
|
|
avg_cost.backward()
|
|
optimizer.minimize(avg_cost)
|
|
transformer.clear_gradients()
|
|
if step_idx % args.print_step == 0:
|
|
total_avg_cost = avg_cost.numpy() * trainer_count
|
|
avg_loss.append(float(total_avg_cost))
|
|
if step_idx == 0:
|
|
logging.info(
|
|
f"step_idx: {step_idx}, epoch: {pass_id}, batch: {batch_id}, avg loss: {total_avg_cost:f}, "
|
|
f"normalized loss: {total_avg_cost - loss_normalizer:f}, ppl: {np.exp([min(total_avg_cost, 100)]).item():f}"
|
|
)
|
|
avg_batch_time = time.time()
|
|
else:
|
|
logging.info(
|
|
f"step_idx: {step_idx}, epoch: {pass_id}, batch: {batch_id}, avg loss: {total_avg_cost:f}, "
|
|
f"normalized loss: {total_avg_cost - loss_normalizer:f}, "
|
|
f"ppl: {np.exp([min(total_avg_cost, 100)]).item():f}, "
|
|
f"speed: {args.print_step / (time.time() - avg_batch_time):.2f} steps/s"
|
|
)
|
|
ce_ppl.append(np.exp([min(total_avg_cost, 100)]))
|
|
avg_batch_time = time.time()
|
|
batch_id += 1
|
|
step_idx += 1
|
|
if step_idx == STEP_NUM:
|
|
if args.save_dygraph_model_path:
|
|
model_dir = os.path.join(args.save_dygraph_model_path)
|
|
if not os.path.exists(model_dir):
|
|
os.makedirs(model_dir)
|
|
paddle.save(
|
|
transformer.state_dict(),
|
|
os.path.join(model_dir, "transformer") + '.pdparams',
|
|
)
|
|
paddle.save(
|
|
optimizer.state_dict(),
|
|
os.path.join(model_dir, "transformer") + '.pdparams',
|
|
)
|
|
break
|
|
time_consumed = time.time() - pass_start_time
|
|
ce_time.append(time_consumed)
|
|
return np.array(avg_loss)
|
|
|
|
|
|
def predict_dygraph(args, batch_generator):
|
|
paddle.seed(SEED)
|
|
paddle.framework.random._manual_program_seed(SEED)
|
|
|
|
# define data loader
|
|
test_loader = paddle.io.DataLoader(
|
|
batch_generator, batch_size=None, places=place
|
|
)
|
|
|
|
# define model
|
|
transformer = paddle.jit.to_static(
|
|
Transformer(
|
|
args.src_vocab_size,
|
|
args.trg_vocab_size,
|
|
args.max_length + 1,
|
|
args.n_layer,
|
|
args.n_head,
|
|
args.d_key,
|
|
args.d_value,
|
|
args.d_model,
|
|
args.d_inner_hid,
|
|
args.prepostprocess_dropout,
|
|
args.attention_dropout,
|
|
args.relu_dropout,
|
|
args.preprocess_cmd,
|
|
args.postprocess_cmd,
|
|
args.weight_sharing,
|
|
args.bos_idx,
|
|
args.eos_idx,
|
|
)
|
|
)
|
|
|
|
# load the trained model
|
|
model_dict, _ = util.load_dygraph(
|
|
os.path.join(args.save_dygraph_model_path, "transformer")
|
|
)
|
|
# to avoid a longer length than training, reset the size of position
|
|
# encoding to max_length
|
|
model_dict["encoder.pos_encoder.weight"] = position_encoding_init(
|
|
args.max_length + 1, args.d_model
|
|
)
|
|
model_dict["decoder.pos_encoder.weight"] = position_encoding_init(
|
|
args.max_length + 1, args.d_model
|
|
)
|
|
transformer.load_dict(model_dict)
|
|
|
|
# set evaluate mode
|
|
transformer.eval()
|
|
|
|
step_idx = 0
|
|
speed_list = []
|
|
for input_data in test_loader():
|
|
(
|
|
src_word,
|
|
src_pos,
|
|
src_slf_attn_bias,
|
|
trg_word,
|
|
trg_src_attn_bias,
|
|
) = input_data
|
|
seq_ids, seq_scores = paddle.jit.to_static(
|
|
transformer.beam_search(
|
|
src_word,
|
|
src_pos,
|
|
src_slf_attn_bias,
|
|
trg_word,
|
|
trg_src_attn_bias,
|
|
bos_id=args.bos_idx,
|
|
eos_id=args.eos_idx,
|
|
beam_size=args.beam_size,
|
|
max_len=args.max_out_len,
|
|
)
|
|
)
|
|
seq_ids = seq_ids.numpy()
|
|
seq_scores = seq_scores.numpy()
|
|
if step_idx % args.print_step == 0:
|
|
if step_idx == 0:
|
|
logging.info(
|
|
f"Dygraph Predict: step_idx: {step_idx}, 1st seq_id: {seq_ids[0][0][0]}, 1st seq_score: {seq_scores[0][0]:.2f}"
|
|
)
|
|
avg_batch_time = time.time()
|
|
else:
|
|
speed = args.print_step / (time.time() - avg_batch_time)
|
|
speed_list.append(speed)
|
|
logging.info(
|
|
f"Dygraph Predict: step_idx: {step_idx}, 1st seq_id: {seq_ids[0][0][0]}, 1st seq_score: {seq_scores[0][0]:.2f}, speed: {speed:.3f} steps/s"
|
|
)
|
|
avg_batch_time = time.time()
|
|
|
|
step_idx += 1
|
|
if step_idx == STEP_NUM:
|
|
break
|
|
logging.info(
|
|
f"Dygraph Predict: avg_speed: {np.mean(speed_list):.4f} steps/s"
|
|
)
|
|
return seq_ids, seq_scores
|
|
|
|
|
|
class TestTransformer(Dy2StTestBase):
|
|
def setUp(self):
|
|
self.temp_dir = tempfile.TemporaryDirectory()
|
|
|
|
def tearDown(self):
|
|
self.temp_dir.cleanup()
|
|
|
|
def prepare(self, mode='train'):
|
|
args = util.ModelHyperParams()
|
|
args.save_dygraph_model_path = os.path.join(
|
|
self.temp_dir.name, args.save_dygraph_model_path
|
|
)
|
|
args.save_static_model_path = os.path.join(
|
|
self.temp_dir.name, args.save_static_model_path
|
|
)
|
|
args.inference_model_dir = os.path.join(
|
|
self.temp_dir.name, args.inference_model_dir
|
|
)
|
|
args.output_file = os.path.join(self.temp_dir.name, args.output_file)
|
|
batch_generator = util.get_feed_data_reader(args, mode)
|
|
if mode == 'train':
|
|
batch_generator = util.TransedWMT16TrainDataSet(
|
|
batch_generator, args.batch_size * (args.epoch + 1)
|
|
)
|
|
else:
|
|
batch_generator = util.TransedWMT16TestDataSet(
|
|
batch_generator, args.batch_size * (args.epoch + 1)
|
|
)
|
|
return args, batch_generator
|
|
|
|
def _test_train(self):
|
|
args, batch_generator = self.prepare(mode='train')
|
|
static_avg_loss = train_dygraph(args, batch_generator)
|
|
with enable_to_static_guard(False):
|
|
dygraph_avg_loss = train_dygraph(args, batch_generator)
|
|
np.testing.assert_allclose(
|
|
static_avg_loss, dygraph_avg_loss, rtol=1e-05
|
|
)
|
|
|
|
def _test_predict(self):
|
|
args, batch_generator = self.prepare(mode='test')
|
|
static_seq_ids, static_scores = predict_dygraph(args, batch_generator)
|
|
with enable_to_static_guard(False):
|
|
dygraph_seq_ids, dygraph_scores = predict_dygraph(
|
|
args, batch_generator
|
|
)
|
|
|
|
np.testing.assert_allclose(static_seq_ids, dygraph_seq_ids, rtol=1e-05)
|
|
np.testing.assert_allclose(static_scores, dygraph_scores, rtol=1e-05)
|
|
|
|
@test_default_mode_only
|
|
def test_check_result(self):
|
|
self._test_train()
|
|
# TODO(zhangliujie) fix predict fail due to precision misalignment
|
|
# self._test_predict()
|
|
|
|
|
|
if __name__ == '__main__':
|
|
unittest.main()
|