test_machine_translation.py 12.7 KB
Newer Older
1
#   Copyright (c) 2018 PaddlePaddle Authors. All Rights Reserved.
D
dzhwinter 已提交
2
#
D
dzhwinter 已提交
3 4 5
# 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
D
dzhwinter 已提交
6
#
D
dzhwinter 已提交
7
#     http://www.apache.org/licenses/LICENSE-2.0
D
dzhwinter 已提交
8
#
D
dzhwinter 已提交
9 10 11 12 13
# 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.
14

Y
Yang Yu 已提交
15
import contextlib
D
dzhwinter 已提交
16

Y
Yan Chunwei 已提交
17
import numpy as np
18
import paddle
19 20 21 22
import paddle.fluid as fluid
import paddle.fluid.framework as framework
import paddle.fluid.layers as pd
from paddle.fluid.executor import Executor
Y
Yang Yu 已提交
23
import unittest
武毅 已提交
24
import os
Y
Yan Chunwei 已提交
25

P
pangyoki 已提交
26 27
paddle.enable_static()

Y
Yan Chunwei 已提交
28 29
dict_size = 30000
source_dict_dim = target_dict_dim = dict_size
Q
Qiao Longfei 已提交
30 31
hidden_dim = 32
word_dim = 16
Q
Qiao Longfei 已提交
32 33
batch_size = 2
max_length = 8
Y
Yan Chunwei 已提交
34 35
topk_size = 50
trg_dic_size = 10000
Q
Qiao Longfei 已提交
36
beam_size = 2
Y
Yan Chunwei 已提交
37

Q
Qiao Longfei 已提交
38 39 40
decoder_size = hidden_dim


Y
Yang Yu 已提交
41
def encoder(is_sparse):
Q
Qiao Longfei 已提交
42
    # encoder
43 44 45 46 47 48 49 50 51
    src_word_id = pd.data(name="src_word_id",
                          shape=[1],
                          dtype='int64',
                          lod_level=1)
    src_embedding = pd.embedding(input=src_word_id,
                                 size=[dict_size, word_dim],
                                 dtype='float32',
                                 is_sparse=is_sparse,
                                 param_attr=fluid.ParamAttr(name='vemb'))
Q
Qiao Longfei 已提交
52

Q
Qiao Longfei 已提交
53 54 55 56 57
    fc1 = pd.fc(input=src_embedding, size=hidden_dim * 4, act='tanh')
    lstm_hidden0, lstm_0 = pd.dynamic_lstm(input=fc1, size=hidden_dim * 4)
    encoder_out = pd.sequence_last_step(input=lstm_hidden0)
    return encoder_out

Q
Qiao Longfei 已提交
58

Y
Yang Yu 已提交
59
def decoder_train(context, is_sparse):
Q
Qiao Longfei 已提交
60
    # decoder
61 62 63 64 65 66 67 68 69
    trg_language_word = pd.data(name="target_language_word",
                                shape=[1],
                                dtype='int64',
                                lod_level=1)
    trg_embedding = pd.embedding(input=trg_language_word,
                                 size=[dict_size, word_dim],
                                 dtype='float32',
                                 is_sparse=is_sparse,
                                 param_attr=fluid.ParamAttr(name='vemb'))
Q
Qiao Longfei 已提交
70

Q
Qiao Longfei 已提交
71
    rnn = pd.DynamicRNN()
Q
Qiao Longfei 已提交
72 73
    with rnn.block():
        current_word = rnn.step_input(trg_embedding)
Q
Qiao Longfei 已提交
74 75
        pre_state = rnn.memory(init=context)
        current_state = pd.fc(input=[current_word, pre_state],
Q
Qiao Longfei 已提交
76 77
                              size=decoder_size,
                              act='tanh')
Q
Qiao Longfei 已提交
78 79 80 81 82 83

        current_score = pd.fc(input=current_state,
                              size=target_dict_dim,
                              act='softmax')
        rnn.update_memory(pre_state, current_state)
        rnn.output(current_score)
Q
Qiao Longfei 已提交
84 85

    return rnn()
Y
Yan Chunwei 已提交
86 87


Y
Yang Yu 已提交
88
def decoder_decode(context, is_sparse):
Q
Qiao Longfei 已提交
89 90
    init_state = context
    array_len = pd.fill_constant(shape=[1], dtype='int64', value=max_length)
Y
Yang Yu 已提交
91
    counter = pd.zeros(shape=[1], dtype='int64', force_cpu=True)
Q
Qiao Longfei 已提交
92 93 94 95 96 97 98 99 100 101

    # fill the first element with init_state
    state_array = pd.create_array('float32')
    pd.array_write(init_state, array=state_array, i=counter)

    # ids, scores as memory
    ids_array = pd.create_array('int64')
    scores_array = pd.create_array('float32')

    init_ids = pd.data(name="init_ids", shape=[1], dtype="int64", lod_level=2)
102 103 104 105
    init_scores = pd.data(name="init_scores",
                          shape=[1],
                          dtype="float32",
                          lod_level=2)
Q
Qiao Longfei 已提交
106 107 108 109 110 111 112 113 114 115 116 117

    pd.array_write(init_ids, array=ids_array, i=counter)
    pd.array_write(init_scores, array=scores_array, i=counter)

    cond = pd.less_than(x=counter, y=array_len)

    while_op = pd.While(cond=cond)
    with while_op.block():
        pre_ids = pd.array_read(array=ids_array, i=counter)
        pre_state = pd.array_read(array=state_array, i=counter)
        pre_score = pd.array_read(array=scores_array, i=counter)

118
        # expand the recursive_sequence_lengths of pre_state to be the same with pre_score
Q
Qiao Longfei 已提交
119 120
        pre_state_expanded = pd.sequence_expand(pre_state, pre_score)

121 122 123 124
        pre_ids_emb = pd.embedding(input=pre_ids,
                                   size=[dict_size, word_dim],
                                   dtype='float32',
                                   is_sparse=is_sparse)
Q
Qiao Longfei 已提交
125 126

        # use rnn unit to update rnn
127
        current_state = pd.fc(input=[pre_state_expanded, pre_ids_emb],
Q
Qiao Longfei 已提交
128 129
                              size=decoder_size,
                              act='tanh')
130
        current_state_with_lod = pd.lod_reset(x=current_state, y=pre_score)
Q
Qiao Longfei 已提交
131
        # use score to do beam search
132
        current_score = pd.fc(input=current_state_with_lod,
Q
Qiao Longfei 已提交
133 134
                              size=target_dict_dim,
                              act='softmax')
135 136
        topk_scores, topk_indices = pd.topk(current_score, k=beam_size)
        # calculate accumulated scores after topk to reduce computation cost
137 138 139 140 141 142 143 144 145 146
        accu_scores = pd.elementwise_add(x=pd.log(topk_scores),
                                         y=pd.reshape(pre_score, shape=[-1]),
                                         axis=0)
        selected_ids, selected_scores = pd.beam_search(pre_ids,
                                                       pre_score,
                                                       topk_indices,
                                                       accu_scores,
                                                       beam_size,
                                                       end_id=10,
                                                       level=0)
Q
Qiao Longfei 已提交
147 148 149 150 151 152 153 154

        pd.increment(x=counter, value=1, in_place=True)

        # update the memories
        pd.array_write(current_state, array=state_array, i=counter)
        pd.array_write(selected_ids, array=ids_array, i=counter)
        pd.array_write(selected_scores, array=scores_array, i=counter)

155 156 157 158 159
        # update the break condition: up to the max length or all candidates of
        # source sentences have ended.
        length_cond = pd.less_than(x=counter, y=array_len)
        finish_cond = pd.logical_not(pd.is_empty(x=selected_ids))
        pd.logical_and(x=length_cond, y=finish_cond, out=cond)
Q
Qiao Longfei 已提交
160 161

    translation_ids, translation_scores = pd.beam_search_decode(
162
        ids=ids_array, scores=scores_array, beam_size=beam_size, end_id=10)
Q
Qiao Longfei 已提交
163 164 165 166 167 168

    # return init_ids, init_scores

    return translation_ids, translation_scores


武毅 已提交
169
def train_main(use_cuda, is_sparse, is_local=True):
Y
Yang Yu 已提交
170 171 172 173 174 175
    if use_cuda and not fluid.core.is_compiled_with_cuda():
        return
    place = fluid.CUDAPlace(0) if use_cuda else fluid.CPUPlace()

    context = encoder(is_sparse)
    rnn_out = decoder_train(context, is_sparse)
176 177 178 179
    label = pd.data(name="target_language_next_word",
                    shape=[1],
                    dtype='int64',
                    lod_level=1)
Q
Qiao Longfei 已提交
180
    cost = pd.cross_entropy(input=rnn_out, label=label)
Y
Yu Yang 已提交
181
    avg_cost = pd.mean(cost)
Q
Qiao Longfei 已提交
182

183 184 185 186
    optimizer = fluid.optimizer.Adagrad(
        learning_rate=1e-4,
        regularization=fluid.regularizer.L2DecayRegularizer(
            regularization_coeff=0.1))
W
Wu Yi 已提交
187
    optimizer.minimize(avg_cost)
Y
Yan Chunwei 已提交
188

189 190 191
    train_data = paddle.batch(paddle.reader.shuffle(
        paddle.dataset.wmt14.train(dict_size), buf_size=1000),
                              batch_size=batch_size)
Y
Yan Chunwei 已提交
192

193 194 195 196
    feed_order = [
        'src_word_id', 'target_language_word', 'target_language_next_word'
    ]

Y
Yan Chunwei 已提交
197 198
    exe = Executor(place)

武毅 已提交
199 200 201
    def train_loop(main_program):
        exe.run(framework.default_startup_program())

202 203 204 205 206
        feed_list = [
            main_program.global_block().var(var_name) for var_name in feed_order
        ]
        feeder = fluid.DataFeeder(feed_list, place)

武毅 已提交
207
        batch_id = 0
208
        for pass_id in range(1):
武毅 已提交
209 210
            for data in train_data():
                outs = exe.run(main_program,
211
                               feed=feeder.feed(data),
武毅 已提交
212 213
                               fetch_list=[avg_cost])
                avg_cost_val = np.array(outs[0])
214 215
                print('pass_id=' + str(pass_id) + ' batch=' + str(batch_id) +
                      " avg_cost=" + str(avg_cost_val))
武毅 已提交
216 217 218 219 220 221 222
                if batch_id > 3:
                    break
                batch_id += 1

    if is_local:
        train_loop(framework.default_main_program())
    else:
G
gongweibao 已提交
223 224
        port = os.getenv("PADDLE_PSERVER_PORT", "6174")
        pserver_ips = os.getenv("PADDLE_PSERVER_IPS")  # ip,ip...
武毅 已提交
225 226 227 228
        eplist = []
        for ip in pserver_ips.split(","):
            eplist.append(':'.join([ip, port]))
        pserver_endpoints = ",".join(eplist)  # ip:port,ip:port...
G
gongweibao 已提交
229
        trainers = int(os.getenv("PADDLE_TRAINERS"))
武毅 已提交
230
        current_endpoint = os.getenv("POD_IP") + ":" + port
G
gongweibao 已提交
231 232
        trainer_id = int(os.getenv("PADDLE_TRAINER_ID"))
        training_role = os.getenv("PADDLE_TRAINING_ROLE", "TRAINER")
武毅 已提交
233
        t = fluid.DistributeTranspiler()
Y
Yancey1989 已提交
234
        t.transpile(trainer_id, pservers=pserver_endpoints, trainers=trainers)
武毅 已提交
235 236 237 238 239 240 241 242
        if training_role == "PSERVER":
            pserver_prog = t.get_pserver_program(current_endpoint)
            pserver_startup = t.get_startup_program(current_endpoint,
                                                    pserver_prog)
            exe.run(pserver_startup)
            exe.run(pserver_prog)
        elif training_role == "TRAINER":
            train_loop(t.get_trainer_program())
Y
Yan Chunwei 已提交
243 244


Y
Yang Yu 已提交
245 246 247 248 249 250 251
def decode_main(use_cuda, is_sparse):
    if use_cuda and not fluid.core.is_compiled_with_cuda():
        return
    place = fluid.CUDAPlace(0) if use_cuda else fluid.CPUPlace()

    context = encoder(is_sparse)
    translation_ids, translation_scores = decoder_decode(context, is_sparse)
Q
Qiao Longfei 已提交
252 253 254 255 256

    exe = Executor(place)
    exe.run(framework.default_startup_program())

    init_ids_data = np.array([1 for _ in range(batch_size)], dtype='int64')
257 258
    init_scores_data = np.array([1. for _ in range(batch_size)],
                                dtype='float32')
Q
Qiao Longfei 已提交
259 260
    init_ids_data = init_ids_data.reshape((batch_size, 1))
    init_scores_data = init_scores_data.reshape((batch_size, 1))
261 262
    init_recursive_seq_lens = [1] * batch_size
    init_recursive_seq_lens = [init_recursive_seq_lens, init_recursive_seq_lens]
Q
Qiao Longfei 已提交
263

264 265 266 267
    init_ids = fluid.create_lod_tensor(init_ids_data, init_recursive_seq_lens,
                                       place)
    init_scores = fluid.create_lod_tensor(init_scores_data,
                                          init_recursive_seq_lens, place)
268

269 270 271
    train_data = paddle.batch(paddle.reader.shuffle(
        paddle.dataset.wmt14.train(dict_size), buf_size=1000),
                              batch_size=batch_size)
Q
Qiao Longfei 已提交
272

273 274 275 276 277 278 279 280
    feed_order = ['src_word_id']
    feed_list = [
        framework.default_main_program().global_block().var(var_name)
        for var_name in feed_order
    ]
    feeder = fluid.DataFeeder(feed_list, place)

    for data in train_data():
281
        feed_dict = feeder.feed([[x[0]] for x in data])
282 283
        feed_dict['init_ids'] = init_ids
        feed_dict['init_scores'] = init_scores
Q
Qiao Longfei 已提交
284 285 286

        result_ids, result_scores = exe.run(
            framework.default_main_program(),
287
            feed=feed_dict,
Q
Qiao Longfei 已提交
288 289
            fetch_list=[translation_ids, translation_scores],
            return_numpy=False)
290
        print(result_ids.recursive_sequence_lengths())
Q
Qiao Longfei 已提交
291 292 293
        break


Y
Yang Yu 已提交
294 295 296 297 298 299 300 301 302 303 304 305 306 307 308
class TestMachineTranslation(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, is_sparse):
309 310
    f_name = 'test_{0}_{1}_train'.format('cuda' if use_cuda else 'cpu',
                                         'sparse' if is_sparse else 'dense')
Y
Yang Yu 已提交
311 312 313 314 315 316 317 318 319

    def f(*args):
        with scope_prog_guard():
            train_main(use_cuda, is_sparse)

    setattr(TestMachineTranslation, f_name, f)


def inject_test_decode(use_cuda, is_sparse, decorator=None):
320 321
    f_name = 'test_{0}_{1}_decode'.format('cuda' if use_cuda else 'cpu',
                                          'sparse' if is_sparse else 'dense')
Y
Yang Yu 已提交
322 323 324 325 326 327 328 329 330 331 332 333 334 335 336 337 338 339 340 341 342 343 344

    def f(*args):
        with scope_prog_guard():
            decode_main(use_cuda, is_sparse)

    if decorator is not None:
        f = decorator(f)

    setattr(TestMachineTranslation, f_name, f)


for _use_cuda_ in (False, True):
    for _is_sparse_ in (False, True):
        inject_test_train(_use_cuda_, _is_sparse_)

for _use_cuda_ in (False, True):
    for _is_sparse_ in (False, True):

        _decorator_ = None
        if _use_cuda_:
            _decorator_ = unittest.skip(
                reason='Beam Search does not support CUDA!')

345 346 347
        inject_test_decode(is_sparse=_is_sparse_,
                           use_cuda=_use_cuda_,
                           decorator=_decorator_)
Y
Yang Yu 已提交
348

Y
Yan Chunwei 已提交
349
if __name__ == '__main__':
Y
Yang Yu 已提交
350
    unittest.main()