# 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 random import unittest import numpy as np import paddle import paddle.nn as nn from paddle import Model, set_device from paddle.static import InputSpec as Input from paddle.fluid.dygraph import Layer from paddle.nn import BeamSearchDecoder, dynamic_decode import paddle.fluid as fluid import paddle.fluid.layers as layers import paddle.fluid.core as core from paddle.fluid.executor import Executor from paddle.fluid.framework import _test_eager_guard paddle.enable_static() class EncoderCell(layers.RNNCell): def __init__(self, num_layers, hidden_size, dropout_prob=0.): self.num_layers = num_layers self.hidden_size = hidden_size self.dropout_prob = dropout_prob self.lstm_cells = [ layers.LSTMCell(hidden_size) for i in range(num_layers) ] def call(self, step_input, states): new_states = [] for i in range(self.num_layers): out, new_state = self.lstm_cells[i](step_input, states[i]) step_input = layers.dropout( out, self.dropout_prob) if self.dropout_prob > 0 else out new_states.append(new_state) return step_input, new_states @property def state_shape(self): return [cell.state_shape for cell in self.lstm_cells] class DecoderCell(layers.RNNCell): def __init__(self, num_layers, hidden_size, dropout_prob=0.): self.num_layers = num_layers self.hidden_size = hidden_size self.dropout_prob = dropout_prob self.lstm_cells = [ layers.LSTMCell(hidden_size) for i in range(num_layers) ] def attention(self, hidden, encoder_output, encoder_padding_mask): query = layers.fc(hidden, size=encoder_output.shape[-1], bias_attr=False) attn_scores = layers.matmul(layers.unsqueeze(query, [1]), encoder_output, transpose_y=True) if encoder_padding_mask is not None: attn_scores = layers.elementwise_add(attn_scores, encoder_padding_mask) attn_scores = layers.softmax(attn_scores) attn_out = layers.squeeze(layers.matmul(attn_scores, encoder_output), [1]) attn_out = layers.concat([attn_out, hidden], 1) attn_out = layers.fc(attn_out, size=self.hidden_size, bias_attr=False) return attn_out def call(self, step_input, states, encoder_output, encoder_padding_mask=None): lstm_states, input_feed = states new_lstm_states = [] step_input = layers.concat([step_input, input_feed], 1) for i in range(self.num_layers): out, new_lstm_state = self.lstm_cells[i](step_input, lstm_states[i]) step_input = layers.dropout( out, self.dropout_prob) if self.dropout_prob > 0 else out new_lstm_states.append(new_lstm_state) out = self.attention(step_input, encoder_output, encoder_padding_mask) return out, [new_lstm_states, out] class Encoder(object): def __init__(self, num_layers, hidden_size, dropout_prob=0.): self.encoder_cell = EncoderCell(num_layers, hidden_size, dropout_prob) def __call__(self, src_emb, src_sequence_length): encoder_output, encoder_final_state = layers.rnn( cell=self.encoder_cell, inputs=src_emb, sequence_length=src_sequence_length, is_reverse=False) return encoder_output, encoder_final_state class Decoder(object): def __init__(self, num_layers, hidden_size, dropout_prob, decoding_strategy="infer_sample", max_decoding_length=20): self.decoder_cell = DecoderCell(num_layers, hidden_size, dropout_prob) self.decoding_strategy = decoding_strategy self.max_decoding_length = None if ( self.decoding_strategy == "train_greedy") else max_decoding_length def __call__(self, decoder_initial_states, encoder_output, encoder_padding_mask, **kwargs): output_layer = kwargs.pop("output_layer", None) if self.decoding_strategy == "train_greedy": # for teach-forcing MLE pre-training helper = layers.TrainingHelper(**kwargs) elif self.decoding_strategy == "infer_sample": helper = layers.SampleEmbeddingHelper(**kwargs) elif self.decoding_strategy == "infer_greedy": helper = layers.GreedyEmbeddingHelper(**kwargs) if self.decoding_strategy == "beam_search": beam_size = kwargs.get("beam_size", 4) encoder_output = layers.BeamSearchDecoder.tile_beam_merge_with_batch( encoder_output, beam_size) encoder_padding_mask = layers.BeamSearchDecoder.tile_beam_merge_with_batch( encoder_padding_mask, beam_size) decoder = layers.BeamSearchDecoder(cell=self.decoder_cell, output_fn=output_layer, **kwargs) else: decoder = layers.BasicDecoder(self.decoder_cell, helper, output_fn=output_layer) (decoder_output, decoder_final_state, dec_seq_lengths) = layers.dynamic_decode( decoder, inits=decoder_initial_states, max_step_num=self.max_decoding_length, encoder_output=encoder_output, encoder_padding_mask=encoder_padding_mask, impute_finished=False # for test coverage if self.decoding_strategy == "beam_search" else True, is_test=True if self.decoding_strategy == "beam_search" else False, return_length=True) return decoder_output, decoder_final_state, dec_seq_lengths class Seq2SeqModel(object): """Seq2Seq model: RNN encoder-decoder with attention""" def __init__(self, num_layers, hidden_size, dropout_prob, src_vocab_size, trg_vocab_size, start_token, end_token, decoding_strategy="infer_sample", max_decoding_length=20, beam_size=4): self.start_token, self.end_token = start_token, end_token self.max_decoding_length, self.beam_size = max_decoding_length, beam_size self.src_embeder = paddle.nn.Embedding( src_vocab_size, hidden_size, weight_attr=fluid.ParamAttr(name="source_embedding")) self.trg_embeder = paddle.nn.Embedding( trg_vocab_size, hidden_size, weight_attr=fluid.ParamAttr(name="target_embedding")) self.encoder = Encoder(num_layers, hidden_size, dropout_prob) self.decoder = Decoder(num_layers, hidden_size, dropout_prob, decoding_strategy, max_decoding_length) self.output_layer = lambda x: layers.fc(x, size=trg_vocab_size, num_flatten_dims=len(x.shape) - 1, param_attr=fluid.ParamAttr(), bias_attr=False) def __call__(self, src, src_length, trg=None, trg_length=None): # encoder encoder_output, encoder_final_state = self.encoder( self.src_embeder(src), src_length) decoder_initial_states = [ encoder_final_state, self.decoder.decoder_cell.get_initial_states( batch_ref=encoder_output, shape=[encoder_output.shape[-1]]) ] src_mask = layers.sequence_mask(src_length, maxlen=layers.shape(src)[1], dtype="float32") encoder_padding_mask = (src_mask - 1.0) * 1e9 encoder_padding_mask = layers.unsqueeze(encoder_padding_mask, [1]) # decoder decoder_kwargs = { "inputs": self.trg_embeder(trg), "sequence_length": trg_length, } if self.decoder.decoding_strategy == "train_greedy" else ( { "embedding_fn": self.trg_embeder, "beam_size": self.beam_size, "start_token": self.start_token, "end_token": self.end_token } if self.decoder.decoding_strategy == "beam_search" else { "embedding_fn": self.trg_embeder, "start_tokens": layers.fill_constant_batch_size_like(input=encoder_output, shape=[-1], dtype=src.dtype, value=self.start_token), "end_token": self.end_token }) decoder_kwargs["output_layer"] = self.output_layer (decoder_output, decoder_final_state, dec_seq_lengths) = self.decoder(decoder_initial_states, encoder_output, encoder_padding_mask, **decoder_kwargs) if self.decoder.decoding_strategy == "beam_search": # for inference return decoder_output logits, samples, sample_length = (decoder_output.cell_outputs, decoder_output.sample_ids, dec_seq_lengths) probs = layers.softmax(logits) return probs, samples, sample_length class PolicyGradient(object): """policy gradient""" def __init__(self, lr=None): self.lr = lr def learn(self, act_prob, action, reward, length=None): """ update policy model self.model with policy gradient algorithm """ self.reward = fluid.layers.py_func(func=reward_func, x=[action, length], out=reward) neg_log_prob = layers.cross_entropy(act_prob, action) cost = neg_log_prob * reward cost = (layers.reduce_sum(cost) / layers.reduce_sum(length) ) if length is not None else layers.reduce_mean(cost) optimizer = fluid.optimizer.Adam(self.lr) optimizer.minimize(cost) return cost def reward_func(samples, sample_length): """toy reward""" def discount_reward(reward, sequence_length, discount=1.): return discount_reward_1d(reward, sequence_length, discount) def discount_reward_1d(reward, sequence_length, discount=1., dtype=None): if sequence_length is None: raise ValueError( 'sequence_length must not be `None` for 1D reward.') reward = np.array(reward) sequence_length = np.array(sequence_length) batch_size = reward.shape[0] max_seq_length = np.max(sequence_length) dtype = dtype or reward.dtype if discount == 1.: dmat = np.ones([batch_size, max_seq_length], dtype=dtype) else: steps = np.tile(np.arange(max_seq_length), [batch_size, 1]) mask = np.asarray(steps < (sequence_length - 1)[:, None], dtype=dtype) # Make each row = [discount, ..., discount, 1, ..., 1] dmat = mask * discount + (1 - mask) dmat = np.cumprod(dmat[:, ::-1], axis=1)[:, ::-1] disc_reward = dmat * reward[:, None] disc_reward = mask_sequences(disc_reward, sequence_length, dtype=dtype) return disc_reward def mask_sequences(sequence, sequence_length, dtype=None, time_major=False): sequence = np.array(sequence) sequence_length = np.array(sequence_length) rank = sequence.ndim if rank < 2: raise ValueError("`sequence` must be 2D or higher order.") batch_size = sequence.shape[0] max_time = sequence.shape[1] dtype = dtype or sequence.dtype if time_major: sequence = np.transpose(sequence, axes=[1, 0, 2]) steps = np.tile(np.arange(max_time), [batch_size, 1]) mask = np.asarray(steps < sequence_length[:, None], dtype=dtype) for _ in range(2, rank): mask = np.expand_dims(mask, -1) sequence = sequence * mask if time_major: sequence = np.transpose(sequence, axes=[1, 0, 2]) return sequence samples = np.array(samples) sample_length = np.array(sample_length) # length reward reward = (5 - np.abs(sample_length - 5)).astype("float32") # repeat punishment to trapped into local minima getting all same words # beam search to get more than one sample may also can avoid this for i in range(reward.shape[0]): reward[i] += -10 if sample_length[i] > 1 and np.all( samples[i][:sample_length[i] - 1] == samples[i][0]) else 0 return discount_reward(reward, sample_length, discount=1.).astype("float32") class MLE(object): """teacher-forcing MLE training""" def __init__(self, lr=None): self.lr = lr def learn(self, probs, label, weight=None, length=None): loss = layers.cross_entropy(input=probs, label=label, soft_label=False) max_seq_len = layers.shape(probs)[1] mask = layers.sequence_mask(length, maxlen=max_seq_len, dtype="float32") loss = loss * mask loss = layers.reduce_mean(loss, dim=[0]) loss = layers.reduce_sum(loss) optimizer = fluid.optimizer.Adam(self.lr) optimizer.minimize(loss) return loss class SeqPGAgent(object): def __init__(self, model_cls, alg_cls=PolicyGradient, model_hparams={}, alg_hparams={}, executor=None, main_program=None, startup_program=None, seed=None): self.main_program = fluid.Program( ) if main_program is None else main_program self.startup_program = fluid.Program( ) if startup_program is None else startup_program if seed is not None: self.main_program.random_seed = seed self.startup_program.random_seed = seed self.build_program(model_cls, alg_cls, model_hparams, alg_hparams) self.executor = executor def build_program(self, model_cls, alg_cls, model_hparams, alg_hparams): with fluid.program_guard(self.main_program, self.startup_program): source = fluid.data(name="src", shape=[None, None], dtype="int64") source_length = fluid.data(name="src_sequence_length", shape=[None], dtype="int64") # only for teacher-forcing MLE training target = fluid.data(name="trg", shape=[None, None], dtype="int64") target_length = fluid.data(name="trg_sequence_length", shape=[None], dtype="int64") label = fluid.data(name="label", shape=[None, None, 1], dtype="int64") self.model = model_cls(**model_hparams) self.alg = alg_cls(**alg_hparams) self.probs, self.samples, self.sample_length = self.model( source, source_length, target, target_length) self.samples.stop_gradient = True self.reward = fluid.data( name="reward", shape=[None, None], # batch_size, seq_len dtype=self.probs.dtype) self.samples.stop_gradient = False self.cost = self.alg.learn(self.probs, self.samples, self.reward, self.sample_length) # to define the same parameters between different programs self.pred_program = self.main_program._prune_with_input( [source.name, source_length.name], [self.probs, self.samples, self.sample_length]) def predict(self, feed_dict): samples, sample_length = self.executor.run( self.pred_program, feed=feed_dict, fetch_list=[self.samples, self.sample_length]) return samples, sample_length def learn(self, feed_dict, fetch_list): results = self.executor.run(self.main_program, feed=feed_dict, fetch_list=fetch_list) return results class TestDynamicDecode(unittest.TestCase): def setUp(self): np.random.seed(123) self.model_hparams = { "num_layers": 2, "hidden_size": 32, "dropout_prob": 0.1, "src_vocab_size": 100, "trg_vocab_size": 100, "start_token": 0, "end_token": 1, "decoding_strategy": "infer_greedy", "max_decoding_length": 10 } self.iter_num = iter_num = 2 self.batch_size = batch_size = 4 src_seq_len = 10 trg_seq_len = 12 self.data = { "src": np.random.randint( 2, self.model_hparams["src_vocab_size"], (iter_num * batch_size, src_seq_len)).astype("int64"), "src_sequence_length": np.random.randint(1, src_seq_len, (iter_num * batch_size, )).astype("int64"), "trg": np.random.randint( 2, self.model_hparams["src_vocab_size"], (iter_num * batch_size, trg_seq_len)).astype("int64"), "trg_sequence_length": np.random.randint(1, trg_seq_len, (iter_num * batch_size, )).astype("int64"), "label": np.random.randint( 2, self.model_hparams["src_vocab_size"], (iter_num * batch_size, trg_seq_len, 1)).astype("int64"), } place = core.CUDAPlace( 0) if core.is_compiled_with_cuda() else core.CPUPlace() self.exe = Executor(place) def test_mle_train(self): paddle.enable_static() self.model_hparams["decoding_strategy"] = "train_greedy" agent = SeqPGAgent(model_cls=Seq2SeqModel, alg_cls=MLE, model_hparams=self.model_hparams, alg_hparams={"lr": 0.001}, executor=self.exe, main_program=fluid.Program(), startup_program=fluid.Program(), seed=123) self.exe.run(agent.startup_program) for iter_idx in range(self.iter_num): reward, cost = agent.learn( { "src": self.data["src"][iter_idx * self.batch_size:(iter_idx + 1) * self.batch_size, :], "src_sequence_length": self.data["src_sequence_length"][iter_idx * self.batch_size: (iter_idx + 1) * self.batch_size], "trg": self.data["trg"][iter_idx * self.batch_size:(iter_idx + 1) * self.batch_size, :], "trg_sequence_length": self.data["trg_sequence_length"][iter_idx * self.batch_size: (iter_idx + 1) * self.batch_size], "label": self.data["label"][iter_idx * self.batch_size:(iter_idx + 1) * self.batch_size] }, fetch_list=[agent.cost, agent.cost]) print("iter_idx: %d, reward: %f, cost: %f" % (iter_idx, reward.mean(), cost)) def test_greedy_train(self): paddle.enable_static() self.model_hparams["decoding_strategy"] = "infer_greedy" agent = SeqPGAgent(model_cls=Seq2SeqModel, alg_cls=PolicyGradient, model_hparams=self.model_hparams, alg_hparams={"lr": 0.001}, executor=self.exe, main_program=fluid.Program(), startup_program=fluid.Program(), seed=123) self.exe.run(agent.startup_program) for iter_idx in range(self.iter_num): reward, cost = agent.learn( { "src": self.data["src"][iter_idx * self.batch_size:(iter_idx + 1) * self.batch_size, :], "src_sequence_length": self.data["src_sequence_length"][iter_idx * self.batch_size: (iter_idx + 1) * self.batch_size] }, fetch_list=[agent.reward, agent.cost]) print("iter_idx: %d, reward: %f, cost: %f" % (iter_idx, reward.mean(), cost)) def test_sample_train(self): paddle.enable_static() self.model_hparams["decoding_strategy"] = "infer_sample" agent = SeqPGAgent(model_cls=Seq2SeqModel, alg_cls=PolicyGradient, model_hparams=self.model_hparams, alg_hparams={"lr": 0.001}, executor=self.exe, main_program=fluid.Program(), startup_program=fluid.Program(), seed=123) self.exe.run(agent.startup_program) for iter_idx in range(self.iter_num): reward, cost = agent.learn( { "src": self.data["src"][iter_idx * self.batch_size:(iter_idx + 1) * self.batch_size, :], "src_sequence_length": self.data["src_sequence_length"][iter_idx * self.batch_size: (iter_idx + 1) * self.batch_size] }, fetch_list=[agent.reward, agent.cost]) print("iter_idx: %d, reward: %f, cost: %f" % (iter_idx, reward.mean(), cost)) def test_beam_search_infer(self): paddle.set_default_dtype("float32") paddle.enable_static() self.model_hparams["decoding_strategy"] = "beam_search" main_program = fluid.Program() startup_program = fluid.Program() with fluid.program_guard(main_program, startup_program): source = fluid.data(name="src", shape=[None, None], dtype="int64") source_length = fluid.data(name="src_sequence_length", shape=[None], dtype="int64") model = Seq2SeqModel(**self.model_hparams) output = model(source, source_length) self.exe.run(startup_program) for iter_idx in range(self.iter_num): trans_ids = self.exe.run( program=main_program, feed={ "src": self.data["src"][iter_idx * self.batch_size:(iter_idx + 1) * self.batch_size, :], "src_sequence_length": self.data["src_sequence_length"][iter_idx * self.batch_size:(iter_idx + 1) * self.batch_size] }, fetch_list=[output])[0] def func_dynamic_basic_decoder(self): paddle.disable_static() src = paddle.to_tensor(np.random.randint(8, size=(8, 4))) src_length = paddle.to_tensor(np.random.randint(8, size=(8))) model = Seq2SeqModel(**self.model_hparams) probs, samples, sample_length = model(src, src_length) paddle.enable_static() def test_dynamic_basic_decoder(self): with _test_eager_guard(): self.func_dynamic_basic_decoder() self.func_dynamic_basic_decoder() class ModuleApiTest(unittest.TestCase): @classmethod def setUpClass(cls): cls._np_rand_state = np.random.get_state() cls._py_rand_state = random.getstate() cls._random_seed = 123 np.random.seed(cls._random_seed) random.seed(cls._random_seed) cls.model_cls = type( cls.__name__ + "Model", (Layer, ), { "__init__": cls.model_init_wrapper(cls.model_init), "forward": cls.model_forward }) @classmethod def tearDownClass(cls): np.random.set_state(cls._np_rand_state) random.setstate(cls._py_rand_state) @staticmethod def model_init_wrapper(func): def __impl__(self, *args, **kwargs): Layer.__init__(self) func(self, *args, **kwargs) return __impl__ @staticmethod def model_init(model, *args, **kwargs): raise NotImplementedError( "model_init acts as `Model.__init__`, thus must implement it") @staticmethod def model_forward(model, *args, **kwargs): return model.module(*args, **kwargs) def make_inputs(self): # TODO(guosheng): add default from `self.inputs` raise NotImplementedError( "model_inputs makes inputs for model, thus must implement it") def setUp(self): """ For the model which wraps the module to be tested: Set input data by `self.inputs` list Set init argument values by `self.attrs` list/dict Set model parameter values by `self.param_states` dict Set expected output data by `self.outputs` list We can create a model instance and run once with these. """ self.inputs = [] self.attrs = {} self.param_states = {} self.outputs = [] def _calc_output(self, place, mode="test", dygraph=True): if dygraph: fluid.enable_dygraph(place) else: fluid.disable_dygraph() gen = paddle.seed(self._random_seed) paddle.framework.random._manual_program_seed(self._random_seed) scope = fluid.core.Scope() with fluid.scope_guard(scope): layer = self.model_cls(**self.attrs) if isinstance( self.attrs, dict) else self.model_cls(*self.attrs) model = Model(layer, inputs=self.make_inputs()) model.prepare() if self.param_states: model.load(self.param_states, optim_state=None) return model.predict_batch(self.inputs) def check_output_with_place(self, place, mode="test"): dygraph_output = self._calc_output(place, mode, dygraph=True) stgraph_output = self._calc_output(place, mode, dygraph=False) expect_output = getattr(self, "outputs", None) for actual_t, expect_t in zip(dygraph_output, stgraph_output): np.testing.assert_allclose(actual_t, expect_t, rtol=1e-05, atol=0) if expect_output: for actual_t, expect_t in zip(dygraph_output, expect_output): np.testing.assert_allclose(actual_t, expect_t, rtol=1e-05, atol=0) def check_output(self): devices = ["CPU", "GPU"] if fluid.is_compiled_with_cuda() else ["CPU"] for device in devices: place = set_device(device) self.check_output_with_place(place) class TestBeamSearch(ModuleApiTest): def setUp(self): paddle.set_default_dtype("float64") shape = (8, 32) self.inputs = [ np.random.random(shape).astype("float64"), np.random.random(shape).astype("float64") ] self.outputs = None self.attrs = { "vocab_size": 100, "embed_dim": 32, "hidden_size": 32, } self.param_states = {} @staticmethod def model_init(self, vocab_size, embed_dim, hidden_size, bos_id=0, eos_id=1, beam_size=4, max_step_num=20): embedder = paddle.fluid.dygraph.Embedding(size=[vocab_size, embed_dim], dtype="float64") output_layer = nn.Linear(hidden_size, vocab_size) cell = nn.LSTMCell(embed_dim, hidden_size) self.max_step_num = max_step_num self.beam_search_decoder = BeamSearchDecoder(cell, start_token=bos_id, end_token=eos_id, beam_size=beam_size, embedding_fn=embedder, output_fn=output_layer) @staticmethod def model_forward(model, init_hidden, init_cell): return dynamic_decode(model.beam_search_decoder, [init_hidden, init_cell], max_step_num=model.max_step_num, impute_finished=True, is_test=True)[0] def make_inputs(self): inputs = [ Input([None, self.inputs[0].shape[-1]], "float64", "init_hidden"), Input([None, self.inputs[1].shape[-1]], "float64", "init_cell"), ] return inputs def func_check_output(self): self.setUp() self.make_inputs() self.check_output() def test_check_output(self): with _test_eager_guard(): self.func_check_output() self.func_check_output() if __name__ == '__main__': unittest.main()