# Copyright (c) 2018 PaddlePaddle Authors. All Rights Reserve. # # 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. """ A simple machine translation demo using beam search decoder. """ import contextlib import numpy as np import paddle import paddle.fluid as fluid import paddle.fluid.framework as framework import paddle.fluid.layers as layers from paddle.fluid.executor import Executor from paddle.fluid.contrib.decoder.beam_search_decoder import BeamSearchDecoder, InitState, StateCell, TrainingDecoder import unittest paddle.enable_static() dict_size = 30000 source_dict_dim = target_dict_dim = dict_size src_dict, trg_dict = paddle.dataset.wmt14.get_dict(dict_size) hidden_dim = 32 word_dim = 32 decoder_size = hidden_dim IS_SPARSE = True batch_size = 2 max_length = 8 topk_size = 50 trg_dic_size = 10000 beam_size = 2 def encoder(): # encoder src_word = layers.data(name="src_word", shape=[1], dtype='int64', lod_level=1) src_embedding = layers.embedding(input=src_word, size=[dict_size, word_dim], dtype='float32', is_sparse=IS_SPARSE) fc1 = layers.fc(input=src_embedding, size=hidden_dim * 4, act='tanh') lstm_hidden0, lstm_0 = layers.dynamic_lstm(input=fc1, size=hidden_dim * 4) encoder_out = layers.sequence_last_step(input=lstm_hidden0) return encoder_out def decoder_state_cell(context): h = InitState(init=context, need_reorder=True) state_cell = StateCell(inputs={'x': None}, states={'h': h}, out_state='h') @state_cell.state_updater def updater(state_cell): current_word = state_cell.get_input('x') prev_h = state_cell.get_state('h') # make sure lod of h heritted from prev_h h = layers.fc(input=[prev_h, current_word], size=decoder_size, act='tanh') state_cell.set_state('h', h) return state_cell def decoder_train(state_cell): # decoder trg_language_word = layers.data(name="target_word", shape=[1], dtype='int64', lod_level=1) trg_embedding = layers.embedding(input=trg_language_word, size=[dict_size, word_dim], dtype='float32', is_sparse=IS_SPARSE) decoder = TrainingDecoder(state_cell) with decoder.block(): current_word = decoder.step_input(trg_embedding) decoder.state_cell.compute_state(inputs={'x': current_word}) current_score = layers.fc(input=decoder.state_cell.get_state('h'), size=target_dict_dim, act='softmax') decoder.state_cell.update_states() decoder.output(current_score) return decoder() def decoder_decode(state_cell): init_ids = layers.data(name="init_ids", shape=[1], dtype="int64", lod_level=2) init_scores = layers.data(name="init_scores", shape=[1], dtype="float32", lod_level=2) decoder = BeamSearchDecoder(state_cell=state_cell, init_ids=init_ids, init_scores=init_scores, target_dict_dim=target_dict_dim, word_dim=word_dim, input_var_dict={}, topk_size=topk_size, sparse_emb=IS_SPARSE, max_len=max_length, beam_size=beam_size, end_id=1, name=None) decoder.decode() translation_ids, translation_scores = decoder() return translation_ids, translation_scores def train_main(use_cuda): if use_cuda and not fluid.core.is_compiled_with_cuda(): return place = fluid.CUDAPlace(0) if use_cuda else fluid.CPUPlace() context = encoder() state_cell = decoder_state_cell(context) rnn_out = decoder_train(state_cell) label = layers.data(name="target_next_word", shape=[1], dtype='int64', lod_level=1) cost = layers.cross_entropy(input=rnn_out, label=label) avg_cost = paddle.mean(x=cost) optimizer = fluid.optimizer.Adagrad(learning_rate=1e-3) optimizer.minimize(avg_cost) train_reader = paddle.batch(paddle.reader.shuffle( paddle.dataset.wmt14.train(dict_size), buf_size=1000), batch_size=batch_size) feed_order = ['src_word', 'target_word', 'target_next_word'] exe = Executor(place) def train_loop(main_program): exe.run(framework.default_startup_program()) feed_list = [ main_program.global_block().var(var_name) for var_name in feed_order ] feeder = fluid.DataFeeder(feed_list, place) for pass_id in range(1): for batch_id, data in enumerate(train_reader()): outs = exe.run(main_program, feed=feeder.feed(data), fetch_list=[avg_cost]) avg_cost_val = np.array(outs[0]) print('pass_id=' + str(pass_id) + ' batch=' + str(batch_id) + " avg_cost=" + str(avg_cost_val)) if batch_id > 3: break train_loop(framework.default_main_program()) def decode_main(use_cuda): if use_cuda and not fluid.core.is_compiled_with_cuda(): return place = fluid.CUDAPlace(0) if use_cuda else fluid.CPUPlace() context = encoder() state_cell = decoder_state_cell(context) translation_ids, translation_scores = decoder_decode(state_cell) exe = Executor(place) exe.run(framework.default_startup_program()) init_ids_data = np.array([0 for _ in range(batch_size)], dtype='int64') init_scores_data = np.array([1. for _ in range(batch_size)], dtype='float32') init_ids_data = init_ids_data.reshape((batch_size, 1)) init_scores_data = init_scores_data.reshape((batch_size, 1)) init_lod = [1] * batch_size init_lod = [init_lod, init_lod] init_ids = fluid.create_lod_tensor(init_ids_data, init_lod, place) init_scores = fluid.create_lod_tensor(init_scores_data, init_lod, place) train_reader = paddle.batch(paddle.reader.shuffle( paddle.dataset.wmt14.train(dict_size), buf_size=1000), batch_size=batch_size) feed_order = ['src_word'] feed_list = [ framework.default_main_program().global_block().var(var_name) for var_name in feed_order ] feeder = fluid.DataFeeder(feed_list, place) data = next(train_reader()) feed_dict = feeder.feed([[x[0]] for x in data]) feed_dict['init_ids'] = init_ids feed_dict['init_scores'] = init_scores result_ids, result_scores = exe.run( framework.default_main_program(), feed=feed_dict, fetch_list=[translation_ids, translation_scores], return_numpy=False) print(result_ids.lod()) class TestBeamSearchDecoder(unittest.TestCase): pass @contextlib.contextmanager def scope_prog_guard(): prog = fluid.Program() startup_prog = fluid.Program() scope = fluid.core.Scope() with fluid.scope_guard(scope): with fluid.program_guard(prog, startup_prog): yield def inject_test_train(use_cuda): f_name = 'test_{0}_train'.format('cuda' if use_cuda else 'cpu') def f(*args): with scope_prog_guard(): train_main(use_cuda) setattr(TestBeamSearchDecoder, f_name, f) def inject_test_decode(use_cuda, decorator=None): f_name = 'test_{0}_decode'.format('cuda' if use_cuda else 'cpu') def f(*args): with scope_prog_guard(): decode_main(use_cuda) if decorator is not None: f = decorator(f) setattr(TestBeamSearchDecoder, f_name, f) for _use_cuda_ in (False, True): inject_test_train(_use_cuda_) for _use_cuda_ in (False, True): _decorator_ = None inject_test_decode(use_cuda=_use_cuda_, decorator=_decorator_) if __name__ == '__main__': unittest.main()