paddlepaddle--paddle/test/dygraph_to_static/test_seq2seq.py

# 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 os
import tempfile
import time
import unittest

import numpy as np
from dygraph_to_static_utils import (
    Dy2StTestBase,
    enable_to_static_guard,
    test_default_mode_only,
)
from seq2seq_dygraph_model import AttentionModel, BaseModel
from seq2seq_utils import Seq2SeqModelHyperParams, get_data_iter

import paddle
from paddle.base.framework import unique_name
from paddle.nn import ClipGradByGlobalNorm

STEP_NUM = 10
PRINT_STEP = 2


def prepare_input(batch):
    src_ids, src_mask, tar_ids, tar_mask = batch
    src_ids = src_ids.reshape((src_ids.shape[0], src_ids.shape[1]))
    in_tar = tar_ids[:, :-1]
    label_tar = tar_ids[:, 1:]

    in_tar = in_tar.reshape((in_tar.shape[0], in_tar.shape[1]))
    label_tar = label_tar.reshape((label_tar.shape[0], label_tar.shape[1], 1))
    inputs = [src_ids, in_tar, label_tar, src_mask, tar_mask]
    return inputs, np.sum(tar_mask)


def train(args, attn_model=False):
    with unique_name.guard():
        paddle.seed(2020)

        if attn_model:
            model = paddle.jit.to_static(
                AttentionModel(
                    args.hidden_size,
                    args.src_vocab_size,
                    args.tar_vocab_size,
                    args.batch_size,
                    num_layers=args.num_layers,
                    init_scale=args.init_scale,
                    dropout=args.dropout,
                )
            )
        else:
            model = paddle.jit.to_static(
                BaseModel(
                    args.hidden_size,
                    args.src_vocab_size,
                    args.tar_vocab_size,
                    args.batch_size,
                    num_layers=args.num_layers,
                    init_scale=args.init_scale,
                    dropout=args.dropout,
                )
            )

        global_norm_clip = ClipGradByGlobalNorm(args.max_grad_norm)
        optimizer = paddle.optimizer.SGD(
            args.learning_rate,
            parameters=model.parameters(),
            grad_clip=global_norm_clip,
        )

        model.train()
        train_data_iter = get_data_iter(args.batch_size)

        batch_times = []
        for batch_id, batch in enumerate(train_data_iter):
            total_loss = 0
            word_count = 0.0
            batch_start_time = time.time()
            input_data_feed, word_num = prepare_input(batch)
            input_data_feed = [
                paddle.to_tensor(np_inp) for np_inp in input_data_feed
            ]
            word_count += word_num
            loss = model(input_data_feed)
            loss.backward()
            optimizer.minimize(loss)
            model.clear_gradients()
            total_loss += loss * args.batch_size
            batch_end_time = time.time()
            batch_time = batch_end_time - batch_start_time
            batch_times.append(batch_time)
            if batch_id % PRINT_STEP == 0:
                print(
                    f"Batch:[{batch_id}]; Time: {batch_time:.5f}s; "
                    f"loss: {loss.numpy():.5f}; "
                    f"total_loss: {total_loss.numpy():.5f}; "
                    f"word num: {word_count:.5f}; "
                    f"ppl: {np.exp(total_loss.numpy() / word_count):.5f}"
                )

            if attn_model:
                # NOTE: Please see code of AttentionModel.
                # Because diff exits if call while_loop in static graph, only run 4 batches to pass the test temporarily.
                if batch_id + 1 >= 4:
                    break
            else:
                if batch_id + 1 >= STEP_NUM:
                    break

        model_path = (
            args.attn_model_path if attn_model else args.base_model_path
        )
        model_dir = os.path.join(model_path)

        if not os.path.exists(model_dir):
            os.makedirs(model_dir)
        paddle.save(model.state_dict(), model_dir + '.pdparams')
        return loss.numpy()


def infer(args, attn_model=False):
    if attn_model:
        model = paddle.jit.to_static(
            AttentionModel(
                args.hidden_size,
                args.src_vocab_size,
                args.tar_vocab_size,
                args.batch_size,
                beam_size=args.beam_size,
                num_layers=args.num_layers,
                init_scale=args.init_scale,
                dropout=0.0,
                mode='beam_search',
            )
        )
    else:
        model = paddle.jit.to_static(
            BaseModel(
                args.hidden_size,
                args.src_vocab_size,
                args.tar_vocab_size,
                args.batch_size,
                beam_size=args.beam_size,
                num_layers=args.num_layers,
                init_scale=args.init_scale,
                dropout=0.0,
                mode='beam_search',
            )
        )

    model_path = args.attn_model_path if attn_model else args.base_model_path
    state_dict = paddle.load(model_path + '.pdparams')
    model.set_dict(state_dict)
    model.eval()
    train_data_iter = get_data_iter(args.batch_size, mode='infer')
    for batch_id, batch in enumerate(train_data_iter):
        input_data_feed, word_num = prepare_input(batch)
        input_data_feed = [
            paddle.to_tensor(np_inp) for np_inp in input_data_feed
        ]
        outputs = paddle.jit.to_static(model.beam_search)(input_data_feed)
        break

    return outputs.numpy()


class TestSeq2seq(Dy2StTestBase):
    def setUp(self):
        self.args = Seq2SeqModelHyperParams
        self.temp_dir = tempfile.TemporaryDirectory()
        self.args.base_model_path = os.path.join(
            self.temp_dir.name, self.args.base_model_path
        )
        self.args.attn_model_path = os.path.join(
            self.temp_dir.name, self.args.attn_model_path
        )
        self.args.reload_model = os.path.join(
            self.temp_dir.name, self.args.reload_model
        )

    def tearDown(self):
        self.temp_dir.cleanup()

    def run_dygraph(self, mode="train", attn_model=False):
        with enable_to_static_guard(False):
            if mode == "train":
                return train(self.args, attn_model)
            else:
                return infer(self.args, attn_model)

    def run_static(self, mode="train", attn_model=False):
        if mode == "train":
            return train(self.args, attn_model)
        else:
            return infer(self.args, attn_model)

    def _test_train(self, attn_model=False):
        dygraph_loss = self.run_dygraph(mode="train", attn_model=attn_model)
        static_loss = self.run_static(mode="train", attn_model=attn_model)
        result = np.allclose(dygraph_loss, static_loss)
        self.assertTrue(
            result,
            msg=f"\ndygraph_loss = {dygraph_loss} \nstatic_loss = {static_loss}",
        )

    def _test_predict(self, attn_model=False):
        pred_dygraph = self.run_dygraph(mode="test", attn_model=attn_model)
        pred_static = self.run_static(mode="test", attn_model=attn_model)
        result = np.allclose(pred_static, pred_dygraph)
        self.assertTrue(
            result,
            msg=f"\npred_dygraph = {pred_dygraph} \npred_static = {pred_static}",
        )

    @test_default_mode_only
    def test_base_model(self):
        self._test_train(attn_model=False)
        self._test_predict(attn_model=False)

    @test_default_mode_only
    def test_attn_model(self):
        self._test_train(attn_model=True)
        # TODO(liym27): add predict
        # self._test_predict(attn_model=True)


if __name__ == '__main__':
    unittest.main()