test_jit_save_load.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.

from __future__ import print_function

import os
import pickle
import shutil
import unittest
import numpy as np
import paddle
from paddle.static import InputSpec
import paddle.fluid as fluid
from paddle.fluid.layers.utils import flatten
from paddle.fluid.dygraph import Linear
from paddle.fluid.dygraph import declarative, ProgramTranslator
from paddle.fluid.dygraph.io import INFER_MODEL_SUFFIX, INFER_PARAMS_SUFFIX, INFER_PARAMS_INFO_SUFFIX
from paddle.fluid import unique_name

BATCH_SIZE = 32
BATCH_NUM = 10
SEED = 10


def random_batch_reader(input_size, label_size):
    def _get_random_inputs_and_labels(input_size, label_size):
        np.random.seed(SEED)
        input = np.random.random(size=input_size).astype('float32')
        label = np.random.random(size=label_size).astype('int64')
        return input, label

    def __reader__():
        for _ in range(BATCH_NUM):
            batch_input, batch_label = _get_random_inputs_and_labels(
                [BATCH_SIZE, input_size], [BATCH_SIZE, label_size])
            yield batch_input, batch_label

    return __reader__


class LinearNet(fluid.dygraph.Layer):
    def __init__(self, in_size, out_size):
        super(LinearNet, self).__init__()
        self._linear = Linear(in_size, out_size)

    @declarative
    def forward(self, x):
        return self._linear(x)


class LinearNetWithInputSpec(fluid.dygraph.Layer):
    def __init__(self, in_size, out_size):
        super(LinearNetWithInputSpec, self).__init__()
        self._linear = Linear(in_size, out_size)

    @declarative(input_spec=[InputSpec(shape=[None, 784], dtype='float32')])
    def forward(self, x):
        return self._linear(x)


class LinearNetNotDeclarative(fluid.dygraph.Layer):
    def __init__(self, in_size, out_size):
        super(LinearNetNotDeclarative, self).__init__()
        self._linear = Linear(in_size, out_size)

    def forward(self, x):
        return self._linear(x)


class LinerNetWithLabel(paddle.nn.Layer):
    def __init__(self, in_size, out_size):
        super(LinerNetWithLabel, self).__init__()
        self._linear = Linear(in_size, out_size)

    @declarative(input_spec=[
        InputSpec(
            shape=[None, 784], dtype='float32', name="image"), InputSpec(
                shape=[None, 1], dtype='int64', name="label")
    ])
    def forward(self, x, label):
        out = self._linear(x)
        loss = fluid.layers.cross_entropy(out, label)
        avg_loss = fluid.layers.mean(loss)
        return out, avg_loss


class LinerNetWithPruneInput(paddle.nn.Layer):
    def __init__(self, in_size, out_size):
        super(LinerNetWithPruneInput, self).__init__()
        self._linear = Linear(in_size, out_size)

    @declarative(input_spec=[
        InputSpec(
            shape=[None, 784], dtype='float32', name="image"), InputSpec(
                shape=[None, 1], dtype='int64', name="label")
    ])
    def forward(self, x, label):
        out = self._linear(x)
        loss = fluid.layers.cross_entropy(out, label)
        avg_loss = fluid.layers.mean(loss)
        return out


class LinerNetWithUselessInput(paddle.nn.Layer):
    def __init__(self, in_size, out_size):
        super(LinerNetWithUselessInput, self).__init__()
        self._linear = Linear(in_size, out_size)

    @declarative(input_spec=[
        InputSpec(
            shape=[None, 784], dtype='float32', name="image"), InputSpec(
                shape=[None, 1], dtype='int64', name="label")
    ])
    def forward(self, x, label):
        out = self._linear(x)
        return out


class LinearNetReturnLoss(fluid.dygraph.Layer):
    def __init__(self, in_size, out_size):
        super(LinearNetReturnLoss, self).__init__()
        self._linear = Linear(in_size, out_size)

    @declarative
    def forward(self, x):
        y = self._linear(x)
        z = self._linear(y)
        loss = fluid.layers.mean(z)
        return z, loss


class LinearNetMultiInput(fluid.dygraph.Layer):
    def __init__(self, in_size, out_size):
        super(LinearNetMultiInput, self).__init__()
        self._linear1 = Linear(in_size, out_size)
        self._linear2 = Linear(in_size, out_size)

    @declarative(input_spec=[
        InputSpec(
            [None, 8], dtype='float32'), InputSpec(
                [None, 8], dtype='float32')
    ])
    def forward(self, x, y):
        x_out = self._linear1(x)
        y_out = self._linear2(y)
        loss = fluid.layers.mean(x_out + y_out)
        return x_out, y_out, loss


class LinearNetMultiInput1(fluid.dygraph.Layer):
    def __init__(self, in_size, out_size):
        super(LinearNetMultiInput1, self).__init__()
        self._linear1 = Linear(in_size, out_size)
        self._linear2 = Linear(in_size, out_size)

    @declarative(input_spec=(InputSpec(
        [None, 8], dtype='float32'), InputSpec(
            [None, 8], dtype='float32')))
    def forward(self, x, y):
        x_out = self._linear1(x)
        y_out = self._linear2(y)
        loss = fluid.layers.mean(x_out + y_out)
        return x_out, y_out, loss


class MultiLoadingLinearNet(fluid.dygraph.Layer):
    def __init__(self, size, model_path):
        super(MultiLoadingLinearNet, self).__init__()
        self._linear = Linear(size, size)
        self._load_linear1 = paddle.jit.load(model_path)
        self._load_linear2 = paddle.jit.load(model_path)

    @declarative
    def forward(self, x):
        tmp1 = self._linear(x)
        tmp2 = self._load_linear1(tmp1)
        tmp3 = self._load_linear2(tmp2)
        y = self._linear(tmp3)
        return y


class LinearNetReturnHidden(fluid.dygraph.Layer):
    def __init__(self, in_size, out_size):
        super(LinearNetReturnHidden, self).__init__()
        self._linear_1 = Linear(in_size, out_size)
        self._linear_2 = Linear(in_size, out_size)

    @declarative
    def forward(self, x):
        y = self._linear_1(x)
        z = self._linear_2(y)
        loss = fluid.layers.mean(z)
        return y, loss


class LinearNetWithNestOut(fluid.dygraph.Layer):
    def __init__(self, in_size, out_size):
        super(LinearNetWithNestOut, self).__init__()
        self._linear_1 = Linear(in_size, out_size)
        self._linear_2 = Linear(in_size, out_size)

    @declarative
    def forward(self, x):
        y = self._linear_1(x)
        z = self._linear_2(y)
        out = y + z
        loss = fluid.layers.mean(out)
        return y, [(z, loss), out]


class LinearNetWithDictInput(paddle.nn.Layer):
    def __init__(self, in_size, out_size):
        super(LinearNetWithDictInput, self).__init__()
        self._linear = Linear(in_size, out_size)

    @paddle.jit.to_static(input_spec=[{
        'img': InputSpec(
            shape=[None, 8], dtype='float32', name='img')
    }, {
        'label': InputSpec(
            shape=[None, 1], dtype='int64', name='label')
    }])
    def forward(self, img, label):
        out = self._linear(img['img'])
        # not return loss to avoid prune output
        loss = paddle.nn.functional.cross_entropy(out, label['label'])
        return out


class LinearNetWithDictInputNoPrune(paddle.nn.Layer):
    def __init__(self, in_size, out_size):
        super(LinearNetWithDictInputNoPrune, self).__init__()
        self._linear = Linear(in_size, out_size)

    def forward(self, img):
        out = self._linear(img['img'] + img['img2'])
        return out


class EmptyLayer(paddle.nn.Layer):
    def __init__(self):
        super(EmptyLayer, self).__init__()

    @paddle.jit.to_static
    def forward(self, x):
        return x


class NoParamLayer(paddle.nn.Layer):
    def __init__(self):
        super(NoParamLayer, self).__init__()

    @paddle.jit.to_static
    def forward(self, x, y):
        return x + y


class LinearNetWithMultiStaticFunc(fluid.dygraph.Layer):
    def __init__(self, in_size, out_size):
        super(LinearNetWithMultiStaticFunc, self).__init__()
        self._linear_0 = Linear(in_size, out_size)
        self._linear_1 = Linear(in_size, out_size)
        self._scale = paddle.to_tensor(9.9)

    @paddle.jit.to_static
    def forward(self, x):
        return self._linear_0(x)

    @paddle.jit.to_static
    def forward_no_param(self, x):
        return x

    @paddle.jit.to_static
    def forward_general(self, x):
        return self._linear_0(x) + self._linear_1(x) * self._scale


def train(layer, input_size=784, label_size=1):
    # create optimizer
    sgd = fluid.optimizer.SGDOptimizer(
        learning_rate=0.01, parameter_list=layer.parameters())
    # create data loader
    train_loader = fluid.io.DataLoader.from_generator(capacity=5)
    train_loader.set_batch_generator(
        random_batch_reader(input_size, label_size))
    # train
    for data in train_loader():
        img, label = data
        label.stop_gradient = True

        cost = layer(img)

        loss = fluid.layers.cross_entropy(cost, label)
        avg_loss = fluid.layers.mean(loss)

        avg_loss.backward()
        sgd.minimize(avg_loss)
        layer.clear_gradients()
    return [img], layer, avg_loss


def train_with_label(layer, input_size=784, label_size=1):
    # create optimizer
    sgd = fluid.optimizer.SGDOptimizer(
        learning_rate=0.01, parameter_list=layer.parameters())
    # create data loader
    train_loader = fluid.io.DataLoader.from_generator(capacity=5)
    train_loader.set_batch_generator(
        random_batch_reader(input_size, label_size))
    # train
    for data in train_loader():
        img, label = data
        label.stop_gradient = True

        out, avg_loss = layer(img, label)

        avg_loss.backward()
        sgd.minimize(avg_loss)
        layer.clear_gradients()
    return out


class TestJitSaveLoad(unittest.TestCase):
    def setUp(self):
        self.model_path = "test_jit_save_load/model"
        # enable dygraph mode
        fluid.enable_dygraph()
        # config seed
        paddle.seed(SEED)
        paddle.framework.random._manual_program_seed(SEED)

    def train_and_save_model(self, model_path=None):
        layer = LinearNet(784, 1)
        example_inputs, layer, _ = train(layer)
        final_model_path = model_path if model_path else self.model_path
        orig_input_types = [type(x) for x in example_inputs]
        paddle.jit.save(
            layer=layer, path=final_model_path, input_spec=example_inputs)
        new_input_types = [type(x) for x in example_inputs]
        self.assertEqual(orig_input_types, new_input_types)
        return layer

    def test_save_load(self):
        # train and save model
        train_layer = self.train_and_save_model()
        # load model
        loaded_layer = paddle.jit.load(self.model_path)
        self.load_and_inference(train_layer, loaded_layer)
        self.load_dygraph_state_dict(train_layer)
        self.load_and_finetune(train_layer, loaded_layer)

    def load_and_inference(self, train_layer, infer_layer):
        train_layer.eval()
        infer_layer.eval()
        # inference & compare
        x = fluid.dygraph.to_variable(
            np.random.random((1, 784)).astype('float32'))
        self.assertTrue(
            np.array_equal(train_layer(x).numpy(), infer_layer(x).numpy()))

    def load_and_finetune(self, train_layer, load_train_layer):
        train_layer.train()
        load_train_layer.train()
        # train & compare
        img0, _, train_loss = train(train_layer)
        img1, _, load_train_loss = train(load_train_layer)
        self.assertTrue(
            np.array_equal(train_loss.numpy(), load_train_loss.numpy()))

    def load_dygraph_state_dict(self, train_layer):
        train_layer.eval()
        # construct new model
        new_layer = LinearNet(784, 1)
        orig_state_dict = new_layer.state_dict()
        load_state_dict = paddle.load(self.model_path)
        for structured_name in orig_state_dict:
            self.assertTrue(structured_name in load_state_dict)
        new_layer.set_state_dict(load_state_dict)
        new_layer.eval()
        # inference & compare
        x = fluid.dygraph.to_variable(
            np.random.random((1, 784)).astype('float32'))
        self.assertTrue(
            np.array_equal(train_layer(x).numpy(), new_layer(x).numpy()))

    def test_load_dygraph_no_path(self):
        model_path = "test_jit_save_load.no_path/model_path"
        with self.assertRaises(ValueError):
            model_dict, _ = fluid.dygraph.load_dygraph(model_path)

    def test_jit_load_no_path(self):
        path = "test_jit_save_load.no_path/model_path"
        with self.assertRaises(ValueError):
            loaded_layer = paddle.jit.load(path)


class TestSaveLoadWithNestOut(unittest.TestCase):
    def setUp(self):
        # enable dygraph mode
        fluid.enable_dygraph()

    def test_nest_output(self):
        x = fluid.dygraph.to_variable(
            np.random.random((4, 8)).astype('float32'))

        net = LinearNetWithNestOut(8, 8)
        dy_outs = flatten(net(x))
        net = declarative(net, input_spec=[InputSpec([None, 8], name='x')])

        model_path = "net_with_nest_out/model"
        paddle.jit.save(net, model_path)

        load_net = paddle.jit.load(model_path)
        load_outs = flatten(load_net(x))

        self.assertTrue(len(dy_outs) == 4)
        for dy_out, load_out in zip(dy_outs, load_outs):
            self.assertTrue(np.allclose(dy_out.numpy(), load_out.numpy()))


class TestSaveLoadWithDictInput(unittest.TestCase):
    def test_dict_input(self):
        # NOTE: This net cannot be executed, it is just 
        # a special case for exporting models in model validation
        # We DO NOT recommend this writing way of Layer
        net = LinearNetWithDictInput(8, 8)
        # net.forward.concrete_program.inputs: 
        # (<__main__.LinearNetWithDictInput object at 0x7f2655298a98>, 
        #  {'img': var img : fluid.VarType.LOD_TENSOR.shape(-1, 8).astype(VarType.FP32)}, 
        #  {'label': var label : fluid.VarType.LOD_TENSOR.shape(-1, 1).astype(VarType.INT64)})
        self.assertEqual(len(net.forward.concrete_program.inputs), 3)

        path = "test_jit_save_load_with_dict_input/model"
        # prune inputs
        paddle.jit.save(
            layer=net,
            path=path,
            input_spec=[{
                'img': InputSpec(
                    shape=[None, 8], dtype='float32', name='img')
            }])

        img = paddle.randn(shape=[4, 8], dtype='float32')
        loaded_net = paddle.jit.load(path)
        loaded_out = loaded_net(img)

        # loaded_net._input_spec():
        # [InputSpec(shape=(-1, 8), dtype=VarType.FP32, name=img)]
        self.assertEqual(len(loaded_net._input_spec()), 1)


class TestSaveLoadWithDictInputNoPrune(unittest.TestCase):
    def test_dict_input(self):
        net = LinearNetWithDictInputNoPrune(8, 8)

        path = "test_jit_save_load_with_dict_input_no_prune/model"
        # prune inputs
        paddle.jit.save(
            layer=net,
            path=path,
            input_spec=[{
                'img': InputSpec(
                    shape=[None, 8], dtype='float32', name='img'),
                'img2': InputSpec(
                    shape=[None, 8], dtype='float32', name='img2')
            }])

        img = paddle.randn(shape=[4, 8], dtype='float32')
        img2 = paddle.randn(shape=[4, 8], dtype='float32')
        loaded_net = paddle.jit.load(path)
        loaded_out = loaded_net(img, img2)

        self.assertEqual(len(loaded_net._input_spec()), 2)


class TestSaveLoadWithInputSpec(unittest.TestCase):
    def setUp(self):
        # enable dygraph mode
        fluid.enable_dygraph()

    def test_with_input_spec(self):
        net = LinearNetReturnLoss(8, 8)
        # set x.shape = [None, 8]
        net.forward = declarative(
            net.forward, input_spec=[InputSpec(
                [None, 8], name='x')])

        model_path = "input_spec.output_spec/model"
        # check inputs and outputs
        self.assertTrue(len(net.forward.inputs) == 1)
        input_x = net.forward.inputs[0]
        self.assertTrue(input_x.shape == (-1, 8))
        self.assertTrue(input_x.name == 'x')

        # 1. prune loss
        output_spec = net.forward.outputs[:1]
        paddle.jit.save(net, model_path, output_spec=output_spec)

        # 2. load to infer
        infer_layer = paddle.jit.load(model_path)
        x = fluid.dygraph.to_variable(
            np.random.random((4, 8)).astype('float32'))
        pred = infer_layer(x)

    def test_multi_in_out(self):
        net = LinearNetMultiInput(8, 8)

        model_path = "multi_inout.output_spec1/model"
        # 1. check inputs and outputs
        self.assertTrue(len(net.forward.inputs) == 2)
        input_x = net.forward.inputs[0]
        input_y = net.forward.inputs[1]
        self.assertTrue(input_x.shape == (-1, 8))
        self.assertTrue(input_y.shape == (-1, 8))

        # 2. prune loss
        output_spec = net.forward.outputs[:2]
        paddle.jit.save(net, model_path, output_spec=output_spec)

        # 3. load to infer
        infer_layer = paddle.jit.load(model_path)
        x = fluid.dygraph.to_variable(
            np.random.random((4, 8)).astype('float32'))
        y = fluid.dygraph.to_variable(
            np.random.random((4, 8)).astype('float32'))
        # 4. predict
        pred_x, pred_y = infer_layer(x, y)

        # 1. prune y and loss
        model_path = "multi_inout.output_spec2/model"
        output_spec = net.forward.outputs[:1]
        paddle.jit.save(net, model_path, [input_x], output_spec=output_spec)
        # 2. load again
        infer_layer2 = paddle.jit.load(model_path)
        # 3. predict
        pred_xx = infer_layer2(x)

        # 4. assert pred_x == pred_xx
        self.assertTrue(np.allclose(pred_x.numpy(), pred_xx.numpy()))

    def test_multi_in_out1(self):
        net = LinearNetMultiInput1(8, 8)

        model_path = "multi_inout1.output_spec1/model"
        # 1. check inputs and outputs
        self.assertTrue(len(net.forward.inputs) == 2)
        input_x = net.forward.inputs[0]
        input_y = net.forward.inputs[1]
        self.assertTrue(input_x.shape == (-1, 8))
        self.assertTrue(input_y.shape == (-1, 8))

        # 2. prune loss
        output_spec = net.forward.outputs[:2]
        paddle.jit.save(net, model_path, output_spec=output_spec)

        # 3. load to infer
        infer_layer = paddle.jit.load(model_path)
        x = fluid.dygraph.to_variable(
            np.random.random((4, 8)).astype('float32'))
        y = fluid.dygraph.to_variable(
            np.random.random((4, 8)).astype('float32'))
        # 4. predict
        pred_x, pred_y = infer_layer(x, y)

        # 1. prune y and loss
        model_path = "multi_inout1.output_spec2/model"
        output_spec = net.forward.outputs[:1]
        paddle.jit.save(net, model_path, (input_x, ), output_spec=output_spec)
        # 2. load again
        infer_layer2 = paddle.jit.load(model_path)
        # 3. predict
        pred_xx = infer_layer2(x)

        # 4. assert pred_x == pred_xx
        self.assertTrue(np.allclose(pred_x.numpy(), pred_xx.numpy()))


class TestJitSaveLoadConfig(unittest.TestCase):
    def setUp(self):
        # enable dygraph mode
        fluid.enable_dygraph()
        # config seed
        paddle.seed(SEED)
        paddle.framework.random._manual_program_seed(SEED)

    def test_output_spec(self):
        train_layer = LinearNetReturnLoss(8, 8)
        adam = fluid.optimizer.AdamOptimizer(
            learning_rate=0.1, parameter_list=train_layer.parameters())
        x = fluid.dygraph.to_variable(
            np.random.random((4, 8)).astype('float32'))
        for i in range(10):
            out, loss = train_layer(x)
            loss.backward()
            adam.minimize(loss)
            train_layer.clear_gradients()

        model_path = "save_load_config.output_spec"
        output_spec = [out]
        paddle.jit.save(
            layer=train_layer,
            path=model_path,
            input_spec=[x],
            output_spec=output_spec)

        train_layer.eval()
        infer_layer = paddle.jit.load(model_path)
        x = fluid.dygraph.to_variable(
            np.random.random((4, 8)).astype('float32'))
        self.assertTrue(
            np.array_equal(train_layer(x)[0].numpy(), infer_layer(x).numpy()))

    def test_save_no_support_config_error(self):
        layer = LinearNet(784, 1)
        path = "no_support_config_test"
        with self.assertRaises(ValueError):
            paddle.jit.save(layer=layer, path=path, model_filename="")

    def test_load_empty_model_filename_error(self):
        path = "error_model_filename_test"
        with self.assertRaises(ValueError):
            paddle.jit.load(path, model_filename="")

    def test_load_empty_params_filename_error(self):
        path = "error_params_filename_test"
        with self.assertRaises(ValueError):
            paddle.jit.load(path, params_filename="")

    def test_load_with_no_support_config(self):
        path = "no_support_config_test"
        with self.assertRaises(ValueError):
            paddle.jit.load(path, separate_params=True)


class TestJitMultipleLoading(unittest.TestCase):
    def setUp(self):
        self.linear_size = 4
        self.model_path = "jit_multi_load/model"
        # enable dygraph mode
        fluid.enable_dygraph()
        # config seed
        paddle.seed(SEED)
        paddle.framework.random._manual_program_seed(SEED)
        # train and save base model
        self.train_and_save_orig_model()

    def train_and_save_orig_model(self):
        layer = LinearNet(self.linear_size, self.linear_size)
        example_inputs, layer, _ = train(layer, self.linear_size, 1)
        paddle.jit.save(
            layer=layer, path=self.model_path, input_spec=example_inputs)

    def test_load_model_retransform_inference(self):
        multi_loaded_layer = MultiLoadingLinearNet(self.linear_size,
                                                   self.model_path)
        state_dict = multi_loaded_layer.state_dict()
        name_set = set()
        for _, var in state_dict.items():
            self.assertTrue(var.name not in name_set)
            name_set.add(var.name)


class TestJitPruneModelAndLoad(unittest.TestCase):
    def setUp(self):
        self.linear_size = 4
        self.model_path = "jit_prune_model_and_load/model"
        # enable dygraph mode
        fluid.enable_dygraph()
        # config seed
        paddle.seed(SEED)
        paddle.framework.random._manual_program_seed(SEED)

    def train_and_save(self):
        train_layer = LinearNetReturnHidden(8, 8)
        adam = fluid.optimizer.AdamOptimizer(
            learning_rate=0.1, parameter_list=train_layer.parameters())
        x = fluid.dygraph.to_variable(
            np.random.random((4, 8)).astype('float32'))
        for i in range(10):
            hidden, loss = train_layer(x)
            loss.backward()
            adam.minimize(loss)
            train_layer.clear_gradients()

        output_spec = [hidden]
        paddle.jit.save(
            layer=train_layer,
            path=self.model_path,
            input_spec=[x],
            output_spec=output_spec)

        return train_layer

    def test_load_pruned_model(self):
        train_layer = self.train_and_save()
        train_layer.eval()

        infer_layer = paddle.jit.load(self.model_path)

        x = fluid.dygraph.to_variable(
            np.random.random((4, 8)).astype('float32'))
        self.assertTrue(
            np.array_equal(train_layer(x)[0].numpy(), infer_layer(x).numpy()))

    def test_load_var_not_in_extra_var_info(self):
        self.train_and_save()

        # chage extra var info
        var_info_path = self.model_path + INFER_PARAMS_INFO_SUFFIX
        with open(var_info_path, 'rb') as f:
            extra_var_info = pickle.load(f)
            extra_var_info.clear()
        with open(var_info_path, 'wb') as f:
            pickle.dump(extra_var_info, f, protocol=2)

        with self.assertRaises(RuntimeError):
            paddle.jit.load(self.model_path)


class TestJitSaveMultiCases(unittest.TestCase):
    def setUp(self):
        # enable dygraph mode
        fluid.enable_dygraph()
        # config seed
        paddle.seed(SEED)
        paddle.framework.random._manual_program_seed(SEED)

    def verify_inference_correctness(self,
                                     layer,
                                     model_path,
                                     with_label_and_loss=False,
                                     with_label=False):
        layer.eval()
        loaded_layer = paddle.jit.load(model_path)
        loaded_layer.eval()
        # inference & compare
        x = paddle.to_tensor(np.random.random((1, 784)).astype('float32'))
        if with_label_and_loss:
            y = paddle.to_tensor(np.random.random((1, 1)).astype('int64'))
            pred, _ = layer(x, y)
            pred = pred.numpy()
        elif with_label:
            y = paddle.to_tensor(np.random.random((1, 1)).astype('int64'))
            pred = layer(x, y)
            pred = pred.numpy()
        else:
            pred = layer(x).numpy()
        loaded_pred = loaded_layer(x).numpy()
        self.assertTrue(
            np.array_equal(pred, loaded_pred),
            msg="Result diff when load and inference:\nlayer result:\n{}\n" \
                "loaded layer result:\n{}".format(pred, loaded_pred))

    def test_no_prune_to_static_after_train(self):
        layer = LinearNet(784, 1)

        train(layer)

        model_path = "test_no_prune_to_static_after_train/model"
        paddle.jit.save(layer, model_path)

        self.verify_inference_correctness(layer, model_path)

    def test_no_prune_to_static_no_train(self):
        layer = LinearNetWithInputSpec(784, 1)

        model_path = "test_no_prune_to_static_no_train/model"
        paddle.jit.save(layer, model_path)

        self.verify_inference_correctness(layer, model_path)

    def test_no_prune_no_to_static_after_train(self):
        layer = LinearNetNotDeclarative(784, 1)

        train(layer)

        model_path = "test_no_prune_no_to_static_after_train/model"
        paddle.jit.save(
            layer,
            model_path,
            input_spec=[InputSpec(
                shape=[None, 784], dtype='float32')])

        self.verify_inference_correctness(layer, model_path)

    def test_no_prune_no_to_static_after_train_with_examples(self):
        layer = LinearNetNotDeclarative(784, 1)

        example_inputs, _, _ = train(layer)

        model_path = "test_no_prune_no_to_static_after_train_with_examples/model"
        paddle.jit.save(layer=layer, path=model_path, input_spec=example_inputs)

        self.verify_inference_correctness(layer, model_path)

    def test_no_prune_no_to_static_no_train(self):
        layer = LinearNetNotDeclarative(784, 1)

        model_path = "test_no_prune_no_to_static_no_train/model"
        paddle.jit.save(
            layer,
            model_path,
            input_spec=[InputSpec(
                shape=[None, 784], dtype='float32')])

        self.verify_inference_correctness(layer, model_path)

    def test_prune_to_static_after_train(self):
        layer = LinerNetWithLabel(784, 1)

        out = train_with_label(layer)

        model_path = "test_prune_to_static_after_train/model"
        paddle.jit.save(
            layer,
            model_path,
            input_spec=[
                InputSpec(
                    shape=[None, 784], dtype='float32', name="image")
            ],
            output_spec=[out])

        self.verify_inference_correctness(
            layer, model_path, with_label_and_loss=True)

    def test_prune_to_static_no_train(self):
        layer = LinerNetWithLabel(784, 1)

        model_path = "test_prune_to_static_no_train/model"
        # TODO: no train, cannot get output_spec var here
        # now only can use index
        output_spec = layer.forward.outputs[:1]
        paddle.jit.save(
            layer,
            model_path,
            input_spec=[
                InputSpec(
                    shape=[None, 784], dtype='float32', name="image")
            ],
            output_spec=output_spec)

        self.verify_inference_correctness(
            layer, model_path, with_label_and_loss=True)

    def test_prune_input_to_static_no_train(self):
        layer = LinerNetWithPruneInput(784, 1)

        model_path = "test_prune_input_to_static_no_train/model"
        paddle.jit.save(
            layer,
            model_path,
            input_spec=[
                InputSpec(
                    shape=[None, 784], dtype='float32', name="image")
            ])

        self.verify_inference_correctness(layer, model_path, with_label=True)

    def test_prune_useless_input_to_static_no_train(self):
        layer = LinerNetWithUselessInput(784, 1)

        model_path = "test_prune_useless_input_to_static_no_train/model"
        paddle.jit.save(
            layer,
            model_path,
            input_spec=[
                InputSpec(
                    shape=[None, 784], dtype='float32', name="image")
            ])

        self.verify_inference_correctness(layer, model_path, with_label=True)

    def test_no_prune_input_spec_name_warning(self):
        layer = LinearNetWithInputSpec(784, 1)

        train(layer)

        model_path = "test_no_prune_input_spec_name_warning/model"
        paddle.jit.save(
            layer,
            model_path,
            input_spec=[InputSpec(
                shape=[None, 784], dtype='float32')])
        paddle.jit.save(
            layer,
            model_path,
            input_spec=[
                InputSpec(
                    shape=[None, 784], dtype='float32', name='feed_input')
            ])

        self.verify_inference_correctness(layer, model_path)

    def test_not_prune_output_spec_name_warning(self):
        layer = LinearNet(784, 1)

        train(layer)

        model_path = "test_not_prune_output_spec_name_warning/model"
        out = paddle.to_tensor(np.random.random((1, 1)).astype('float'))
        paddle.jit.save(layer, model_path, output_spec=[out])

        self.verify_inference_correctness(layer, model_path)

    def test_prune_input_spec_name_error(self):
        layer = LinerNetWithLabel(784, 1)

        model_path = "test_prune_input_spec_name_error/model"
        with self.assertRaises(ValueError):
            paddle.jit.save(
                layer,
                model_path,
                input_spec=[InputSpec(
                    shape=[None, 784], dtype='float32')])
        with self.assertRaises(ValueError):
            paddle.jit.save(
                layer,
                model_path,
                input_spec=[
                    InputSpec(
                        shape=[None, 784], dtype='float32', name='feed_input')
                ])

    def test_prune_output_spec_name_error(self):
        layer = LinerNetWithLabel(784, 1)

        train_with_label(layer)

        model_path = "test_prune_to_static_after_train/model"
        out = paddle.to_tensor(np.random.random((1, 1)).astype('float'))
        with self.assertRaises(ValueError):
            paddle.jit.save(
                layer,
                model_path,
                input_spec=[
                    InputSpec(
                        shape=[None, 784], dtype='float32', name="image")
                ],
                output_spec=[out])


class TestJitSaveLoadEmptyLayer(unittest.TestCase):
    def setUp(self):
        self.model_path = "jit_save_load_empty_layer/model"
        # enable dygraph mode
        paddle.disable_static()

    def test_save_load_empty_layer(self):
        layer = EmptyLayer()
        x = paddle.to_tensor(np.random.random((10)).astype('float32'))
        out = layer(x)
        paddle.jit.save(layer, self.model_path)
        load_layer = paddle.jit.load(self.model_path)
        load_out = load_layer(x)
        self.assertTrue(np.array_equal(out, load_out))


class TestJitSaveLoadNoParamLayer(unittest.TestCase):
    def setUp(self):
        self.model_path = "jit_save_load_no_param_layer/model"
        # enable dygraph mode
        paddle.disable_static()

    def test_save_load_no_param_layer(self):
        layer = NoParamLayer()
        x = paddle.to_tensor(np.random.random((5)).astype('float32'))
        y = paddle.to_tensor(np.random.random((5)).astype('float32'))
        out = layer(x, y)
        paddle.jit.save(layer, self.model_path)
        load_layer = paddle.jit.load(self.model_path)
        load_out = load_layer(x, y)
        self.assertTrue(np.array_equal(out, load_out))


class TestJitSaveLoadMultiMethods(unittest.TestCase):
    def setUp(self):
        # enable dygraph mode
        paddle.disable_static()

    def test_jit_save_load_inference(self):
        model_path_inference = "jit_save_load_multi_methods/model"
        IMAGE_SIZE = 224
        layer = LinearNetWithMultiStaticFunc(IMAGE_SIZE, 10)
        inps = paddle.randn([1, IMAGE_SIZE])
        result_origin = {}
        for func in dir(layer):
            if func.startswith('forward'):
                result_origin[func] = getattr(layer, func, None)(inps)
        paddle.jit.save(layer, model_path_inference)
        load_net = paddle.jit.load(model_path_inference)
        for func, result in result_origin.items():
            self.assertTrue(
                float((result - getattr(load_net, func, None)(inps)).abs().max(
                )) < 1e-5)

    def test_jit_save_load_multi_methods_inputspec(self):
        model_path = 'jit_save_load_multi_methods/model'
        layer = LinearNetWithMultiStaticFunc(784, 1)
        with self.assertRaises(ValueError):
            paddle.jit.save(
                layer, model_path, input_spec=[InputSpec(shape=[None, 784])])

    def test_parse_name(self):
        model_path_inference = "jit_save_load_parse_name/model"
        IMAGE_SIZE = 224
        layer = LinearNet(IMAGE_SIZE, 1)
        inps = paddle.randn([1, IMAGE_SIZE])
        layer(inps)
        paddle.jit.save(layer, model_path_inference)
        paddle.jit.save(layer, model_path_inference + '_v2')
        load_net = paddle.jit.load(model_path_inference)

        self.assertFalse(hasattr(load_net, 'v2'))


class LayerSaved(paddle.nn.Layer):
    def __init__(self, in_size, out_size):
        super(LayerSaved, self).__init__()
        self.hidden = 100
        self._linear_0 = Linear(in_size, self.hidden)
        self._linear_1_0 = Linear(self.hidden, self.hidden)
        self._linear_1_1 = Linear(self.hidden, self.hidden)
        self._linear_2 = Linear(self.hidden, out_size)
        self._scale = paddle.to_tensor(9.9)

    @paddle.jit.to_static
    def forward(self, x):
        y = self._linear_0(x)
        # Multiple blocks
        if x.shape[0] == 1:
            y = self._linear_1_0(y)
        else:
            y += self._linear_1_1(y + self._scale)
        return self._linear_2(y)


class LayerLoadFinetune(paddle.nn.Layer):
    def __init__(self, in_size, out_size, load_path):
        super(LayerLoadFinetune, self).__init__()
        # Test duplicate name
        self._linear_0 = Linear(in_size, in_size)
        self._linear_1_0 = Linear(out_size, in_size)
        self._linear_1_1 = Linear(out_size, in_size)
        self._linear_2 = Linear(out_size, out_size)
        self._scale = paddle.to_tensor(9.9)

        # Load multiple times
        self._load_l1 = paddle.jit.load(load_path)
        self._load_l2 = paddle.jit.load(load_path)

    @paddle.jit.to_static
    def forward(self, x):
        y = self._linear_0(x)
        y = self._load_l1(y)
        # Multiple blocks
        if x.shape[0] == 1:
            y = self._linear_1_0(y)
            y = self._load_l1(y)
        else:
            y += self._linear_1_1(x + self._scale)
            y = self._load_l2(y)
        y = self._linear_1_0(y)
        y = self._load_l1(y)
        y = self._linear_1_0(y)
        # Use the same layer multiple times.
        y = self._load_l1(y)
        return y


class TestJitSaveLoadSaveWithoutRunning(unittest.TestCase):
    def setUp(self):
        # enable dygraph mode
        paddle.disable_static()

    def test_save_load_finetune_load(self):
        model_path = "test_jit_save_load_save_without_running/model"
        IMAGE_SIZE = 224
        inps0 = paddle.randn([1, IMAGE_SIZE])
        inps1 = paddle.randn([2, IMAGE_SIZE])
        # Use new namespace
        with unique_name.guard():
            layer_save = LayerSaved(IMAGE_SIZE, IMAGE_SIZE)
        #save
        paddle.jit.save(
            layer_save,
            model_path,
            input_spec=[
                paddle.static.InputSpec(
                    shape=[None, IMAGE_SIZE], dtype='float32')
            ])

        result_00 = layer_save(inps0)
        result_01 = layer_save(inps1)
        #load and save without running
        with unique_name.guard():
            layer_load = paddle.jit.load(model_path)
            paddle.jit.save(
                layer_load,
                model_path,
                input_spec=[
                    paddle.static.InputSpec(
                        shape=[None, IMAGE_SIZE], dtype='float32')
                ])
        #reload
        layer_reload = paddle.jit.load(model_path)
        result_10 = layer_reload(inps0)
        result_11 = layer_reload(inps1)

        self.assertTrue(float((result_00 - result_10).abs().max()) < 1e-5)
        self.assertTrue(float((result_01 - result_11).abs().max()) < 1e-5)


class TestJitSaveLoadFinetuneLoad(unittest.TestCase):
    def setUp(self):
        # enable dygraph mode
        paddle.disable_static()

    def test_save_load_finetune_load(self):
        model_path = "test_jit_save_load_finetune_load/model"
        IMAGE_SIZE = 224
        inps0 = paddle.randn([1, IMAGE_SIZE])
        inps1 = paddle.randn([2, IMAGE_SIZE])
        # Use new namespace
        with unique_name.guard():
            layer_save = LayerSaved(IMAGE_SIZE, IMAGE_SIZE)
        layer_save(inps0)
        #save
        paddle.jit.save(layer_save, model_path)
        #load
        with unique_name.guard():
            layer_load = LayerLoadFinetune(IMAGE_SIZE, IMAGE_SIZE, model_path)
        #train
        train(layer_load, input_size=IMAGE_SIZE)
        result_00 = layer_load(inps0)
        result_01 = layer_load(inps1)
        #save
        paddle.jit.save(layer_load, model_path)
        #load
        layer_finetune = paddle.jit.load(model_path)
        result_10 = layer_finetune(inps0)
        result_11 = layer_finetune(inps1)

        self.assertTrue(float((result_00 - result_10).abs().max()) < 1e-5)
        self.assertTrue(float(((result_01 - result_11)).abs().max()) < 1e-5)


# NOTE(weixin): When there are multiple test functions in an 
# `unittest.TestCase`, functions will affect each other, 
# and there is a risk of random failure. 
# So divided into three TestCase: TestJitSaveLoadFunctionCase1, 
# TestJitSaveLoadFunctionCase2, TestJitSaveLoadFunctionCase3.
class TestJitSaveLoadFunctionCase1(unittest.TestCase):
    def setUp(self):
        paddle.disable_static()

    def test_jit_save_load_static_function(self):
        @paddle.jit.to_static
        def fun(inputs):
            return paddle.tanh(inputs)

        path = 'test_jit_save_load_function_1/func'
        inps = paddle.rand([3, 6])
        origin = fun(inps)

        paddle.jit.save(fun, path)
        load_func = paddle.jit.load(path)

        load_result = load_func(inps)
        self.assertTrue((load_result - origin).abs().max() < 1e-10)


class TestJitSaveLoadFunctionCase2(unittest.TestCase):
    def setUp(self):
        paddle.disable_static()

    def test_jit_save_load_function_input_spec(self):
        @paddle.jit.to_static(input_spec=[
            InputSpec(
                shape=[None, 6], dtype='float32', name='x'),
        ])
        def fun(inputs):
            return paddle.nn.functional.relu(inputs)

        path = 'test_jit_save_load_function_2/func'
        inps = paddle.rand([3, 6])
        origin = fun(inps)

        paddle.jit.save(fun, path)
        load_func = paddle.jit.load(path)
        load_result = load_func(inps)
        self.assertTrue((load_result - origin).abs().max() < 1e-10)


class TestJitSaveLoadFunctionCase3(unittest.TestCase):
    def setUp(self):
        paddle.disable_static()

    def test_jit_save_load_function_function(self):
        def fun(inputs):
            return paddle.tanh(inputs)

        path = 'test_jit_save_load_function_3/func'
        inps = paddle.rand([3, 6])
        origin = fun(inps)

        paddle.jit.save(
            fun,
            path,
            input_spec=[
                InputSpec(
                    shape=[None, 6], dtype='float32', name='x'),
            ])
        load_func = paddle.jit.load(path)

        load_result = load_func(inps)
        self.assertTrue((load_result - origin).abs().max() < 1e-10)


class TestJitSaveLoadDataParallel(unittest.TestCase):
    def verify_inference_correctness(self, layer, path):
        layer.eval()
        loaded_layer = paddle.jit.load(path)
        loaded_layer.eval()
        # inference & compare
        x = paddle.to_tensor(np.random.random((1, 784)).astype('float32'))
        pred = layer(x).numpy()
        loaded_pred = loaded_layer(x).numpy()
        self.assertTrue(
            np.array_equal(pred, loaded_pred),
            msg="Result diff when load and inference:\nlayer result:\n{}\n" \
                "loaded layer result:\n{}".format(pred, loaded_pred))

    def test_jit_save_data_parallel_with_inputspec(self):
        layer = LinearNetNotDeclarative(784, 1)
        layer = paddle.DataParallel(layer)

        path = "jit_save_data_parallel_with_inputspec/model"
        paddle.jit.save(
            layer=layer, path=path, input_spec=[InputSpec(shape=[None, 784])])

        self.verify_inference_correctness(layer, path)

    def test_jit_save_data_parallel_with_to_static(self):
        layer = LinearNetWithInputSpec(784, 1)
        layer = paddle.DataParallel(layer)

        path = "jit_save_data_parallel_with_to_static/model"
        paddle.jit.save(layer, path)

        self.verify_inference_correctness(layer, path)


class InputSepcLayer(paddle.nn.Layer):
    '''
    A layer with InputSpec to test InputSpec compatibility
    '''

    @paddle.jit.to_static(input_spec=[
        InputSpec(
            shape=[None, 8], dtype='float32', name='x'), InputSpec(
                shape=[None, 1], dtype='float64', name='y')
    ])
    def forward(self, x, y):
        return x, y


class TestInputSpecCompatibility(unittest.TestCase):
    def _assert_input_spec_layer_return(self, expect_layer, test_layer):
        input_x = paddle.uniform([8, 8], dtype='float32')
        input_y = paddle.uniform([8, 1], dtype='float64')
        expected_result = expect_layer(input_x, input_y)
        test_result = test_layer(input_x, input_y)
        np.testing.assert_allclose(expected_result[0].numpy(),
                                   test_result[0].numpy())
        np.testing.assert_allclose(expected_result[1].numpy(),
                                   test_result[1].numpy())

    def test_jit_save_compatible_input_sepc(self):
        layer = InputSepcLayer()
        save_dir = "jit_save_compatible_input_spec"
        path = save_dir + "/model"

        paddle.jit.save(layer=layer, path=path)
        no_input_spec_layer = paddle.jit.load(path)
        self._assert_input_spec_layer_return(layer, no_input_spec_layer)
        shutil.rmtree(save_dir)

        paddle.jit.save(
            layer=layer,
            path=path,
            input_spec=[
                InputSpec(
                    shape=[None, 8], dtype='float32', name='x'), InputSpec(
                        shape=[None, 1], dtype='float64', name='y')
            ])
        same_input_spec_layer = paddle.jit.load(path)
        self._assert_input_spec_layer_return(layer, same_input_spec_layer)
        shutil.rmtree(save_dir)

        paddle.jit.save(
            layer=layer,
            path=path,
            input_spec=[
                InputSpec(
                    shape=[8, 8], dtype='float32'), InputSpec(
                        shape=[8, -1], dtype='float64')
            ])
        compatible_input_spec_layer = paddle.jit.load(path)
        self._assert_input_spec_layer_return(layer, compatible_input_spec_layer)
        shutil.rmtree(save_dir)

    def test_jit_save_incompatible_input_sepc(self):
        layer = InputSepcLayer()
        save_dir = "jit_save_compatible_input_spec"
        path = save_dir + "/model"

        with self.assertRaises(ValueError):
            # type mismatch
            paddle.jit.save(
                layer=layer,
                path=path,
                input_spec=[
                    InputSpec(
                        shape=[None, 8], dtype='float64'), InputSpec(
                            shape=[None, 1], dtype='float64')
                ])

        with self.assertRaises(ValueError):
            # shape len mismatch
            paddle.jit.save(
                layer=layer,
                path=path,
                input_spec=[
                    InputSpec(
                        shape=[None, 8, 1], dtype='float32'), InputSpec(
                            shape=[None, 1], dtype='float64')
                ])

        with self.assertRaises(ValueError):
            # shape mismatch
            paddle.jit.save(
                layer=layer,
                path=path,
                input_spec=[
                    InputSpec(
                        shape=[None, 8], dtype='float32'), InputSpec(
                            shape=[None, 2], dtype='float64')
                ])
        if os.path.exists(save_dir):
            shutil.rmtree(save_dir)


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