test_jit_save_load.py 55.1 KB
Newer Older
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16
# 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

17
import os
18
import pickle
19
import shutil
20
import unittest
21
import tempfile
22
import numpy as np
L
Leo Chen 已提交
23
import paddle
24
from paddle.static import InputSpec
25
import paddle.fluid as fluid
26
from paddle.fluid.layers.utils import flatten
27
from paddle.fluid.dygraph import Linear
28
from paddle.fluid.dygraph import declarative, ProgramTranslator
29
from paddle.fluid.dygraph.io import INFER_MODEL_SUFFIX, INFER_PARAMS_SUFFIX, INFER_PARAMS_INFO_SUFFIX
W
WeiXin 已提交
30
from paddle.fluid import unique_name
31 32

BATCH_SIZE = 32
33
BATCH_NUM = 10
34 35 36
SEED = 10


37
def random_batch_reader(input_size, label_size):
38

39
    def _get_random_inputs_and_labels(input_size, label_size):
40
        np.random.seed(SEED)
41 42 43
        input = np.random.random(size=input_size).astype('float32')
        label = np.random.random(size=label_size).astype('int64')
        return input, label
44 45 46

    def __reader__():
        for _ in range(BATCH_NUM):
47 48 49
            batch_input, batch_label = _get_random_inputs_and_labels(
                [BATCH_SIZE, input_size], [BATCH_SIZE, label_size])
            yield batch_input, batch_label
50 51 52 53 54

    return __reader__


class LinearNet(fluid.dygraph.Layer):
55

56 57 58 59 60 61 62 63 64
    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)


65
class LinearNetWithInputSpec(fluid.dygraph.Layer):
66

67 68 69 70 71 72 73 74 75
    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)


76
class LinearNetNotDeclarative(fluid.dygraph.Layer):
77

78 79 80 81 82 83 84 85
    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)


86
class LinerNetWithLabel(paddle.nn.Layer):
87

88 89 90 91 92
    def __init__(self, in_size, out_size):
        super(LinerNetWithLabel, self).__init__()
        self._linear = Linear(in_size, out_size)

    @declarative(input_spec=[
93 94
        InputSpec(shape=[None, 784], dtype='float32', name="image"),
        InputSpec(shape=[None, 1], dtype='int64', name="label")
95 96 97 98 99 100 101 102
    ])
    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


C
Chen Weihang 已提交
103
class LinerNetWithPruneInput(paddle.nn.Layer):
104

C
Chen Weihang 已提交
105 106 107 108 109
    def __init__(self, in_size, out_size):
        super(LinerNetWithPruneInput, self).__init__()
        self._linear = Linear(in_size, out_size)

    @declarative(input_spec=[
110 111
        InputSpec(shape=[None, 784], dtype='float32', name="image"),
        InputSpec(shape=[None, 1], dtype='int64', name="label")
C
Chen Weihang 已提交
112 113 114 115 116 117 118 119 120
    ])
    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):
121

C
Chen Weihang 已提交
122 123 124 125 126
    def __init__(self, in_size, out_size):
        super(LinerNetWithUselessInput, self).__init__()
        self._linear = Linear(in_size, out_size)

    @declarative(input_spec=[
127 128
        InputSpec(shape=[None, 784], dtype='float32', name="image"),
        InputSpec(shape=[None, 1], dtype='int64', name="label")
C
Chen Weihang 已提交
129 130 131 132 133 134
    ])
    def forward(self, x, label):
        out = self._linear(x)
        return out


135
class LinearNetReturnLoss(fluid.dygraph.Layer):
136

137 138 139 140 141 142 143 144 145 146 147 148
    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


149
class LinearNetMultiInput(fluid.dygraph.Layer):
150

151 152 153 154 155 156
    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=[
157 158
        InputSpec([None, 8], dtype='float32'),
        InputSpec([None, 8], dtype='float32')
159 160 161 162 163 164 165 166
    ])
    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


167
class LinearNetMultiInput1(fluid.dygraph.Layer):
168

169 170 171 172 173
    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)

174 175
    @declarative(input_spec=(InputSpec([None, 8], dtype='float32'),
                             InputSpec([None, 8], dtype='float32')))
176 177 178 179 180 181 182
    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


183
class MultiLoadingLinearNet(fluid.dygraph.Layer):
184

185 186 187
    def __init__(self, size, model_path):
        super(MultiLoadingLinearNet, self).__init__()
        self._linear = Linear(size, size)
188 189
        self._load_linear1 = paddle.jit.load(model_path)
        self._load_linear2 = paddle.jit.load(model_path)
190 191 192 193 194 195 196 197 198 199 200

    @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):
201

202 203 204 205 206 207 208 209 210 211 212 213 214
    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


215
class LinearNetWithNestOut(fluid.dygraph.Layer):
216

217 218 219 220 221 222 223 224 225 226 227 228 229 230
    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]


231
class LinearNetWithDictInput(paddle.nn.Layer):
232

233 234 235 236 237
    def __init__(self, in_size, out_size):
        super(LinearNetWithDictInput, self).__init__()
        self._linear = Linear(in_size, out_size)

    @paddle.jit.to_static(input_spec=[{
238 239
        'img':
        InputSpec(shape=[None, 8], dtype='float32', name='img')
240
    }, {
241 242
        'label':
        InputSpec(shape=[None, 1], dtype='int64', name='label')
243 244 245 246 247 248 249 250
    }])
    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


251
class LinearNetWithDictInputNoPrune(paddle.nn.Layer):
252

253 254 255 256 257 258 259 260 261
    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


262
class EmptyLayer(paddle.nn.Layer):
263

264 265 266 267 268 269 270 271 272
    def __init__(self):
        super(EmptyLayer, self).__init__()

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


class NoParamLayer(paddle.nn.Layer):
273

274 275 276 277 278 279 280 281
    def __init__(self):
        super(NoParamLayer, self).__init__()

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


282
class LinearNetWithMultiStaticFunc(fluid.dygraph.Layer):
283

284 285 286 287 288 289 290 291 292 293 294 295 296 297 298 299 300 301 302
    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


303
def train(layer, input_size=784, label_size=1):
304
    # create optimizer
305 306
    sgd = fluid.optimizer.SGDOptimizer(learning_rate=0.01,
                                       parameter_list=layer.parameters())
307 308
    # create data loader
    train_loader = fluid.io.DataLoader.from_generator(capacity=5)
309 310
    train_loader.set_batch_generator(random_batch_reader(
        input_size, label_size))
311 312 313 314 315 316 317 318 319 320 321
    # 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()
L
Leo Chen 已提交
322
        sgd.minimize(avg_loss)
323 324 325 326
        layer.clear_gradients()
    return [img], layer, avg_loss


327 328
def train_with_label(layer, input_size=784, label_size=1):
    # create optimizer
329 330
    sgd = fluid.optimizer.SGDOptimizer(learning_rate=0.01,
                                       parameter_list=layer.parameters())
331 332
    # create data loader
    train_loader = fluid.io.DataLoader.from_generator(capacity=5)
333 334
    train_loader.set_batch_generator(random_batch_reader(
        input_size, label_size))
335 336 337 338 339 340 341 342 343 344 345 346 347
    # 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


348
class TestJitSaveLoad(unittest.TestCase):
349

350
    def setUp(self):
351 352 353
        self.temp_dir = tempfile.TemporaryDirectory()
        self.model_path = os.path.join(self.temp_dir.name,
                                       "test_jit_save_load/model")
354 355 356
        # enable dygraph mode
        fluid.enable_dygraph()
        # config seed
C
cnn 已提交
357
        paddle.seed(SEED)
L
Leo Chen 已提交
358
        paddle.framework.random._manual_program_seed(SEED)
359

360 361 362
    def tearDown(self):
        self.temp_dir.cleanup()

363
    def train_and_save_model(self, model_path=None):
364 365
        layer = LinearNet(784, 1)
        example_inputs, layer, _ = train(layer)
366
        final_model_path = model_path if model_path else self.model_path
367
        orig_input_types = [type(x) for x in example_inputs]
368 369 370
        paddle.jit.save(layer=layer,
                        path=final_model_path,
                        input_spec=example_inputs)
371 372
        new_input_types = [type(x) for x in example_inputs]
        self.assertEqual(orig_input_types, new_input_types)
373 374
        return layer

375
    def test_save_load(self):
376 377 378
        # train and save model
        train_layer = self.train_and_save_model()
        # load model
379
        loaded_layer = paddle.jit.load(self.model_path)
380 381 382 383 384
        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):
385
        train_layer.eval()
386
        infer_layer.eval()
387 388 389 390
        # inference & compare
        x = fluid.dygraph.to_variable(
            np.random.random((1, 784)).astype('float32'))
        self.assertTrue(
391 392
            np.array_equal(train_layer(x).numpy(),
                           infer_layer(x).numpy()))
393

394 395
    def load_and_finetune(self, train_layer, load_train_layer):
        train_layer.train()
396 397
        load_train_layer.train()
        # train & compare
L
Leo Chen 已提交
398 399
        img0, _, train_loss = train(train_layer)
        img1, _, load_train_loss = train(load_train_layer)
400 401 402
        self.assertTrue(
            np.array_equal(train_loss.numpy(), load_train_loss.numpy()))

403 404
    def load_dygraph_state_dict(self, train_layer):
        train_layer.eval()
405
        # construct new model
406
        new_layer = LinearNet(784, 1)
407
        orig_state_dict = new_layer.state_dict()
408
        load_state_dict = paddle.load(self.model_path)
409 410 411
        for structured_name in orig_state_dict:
            self.assertTrue(structured_name in load_state_dict)
        new_layer.set_state_dict(load_state_dict)
412 413 414 415 416
        new_layer.eval()
        # inference & compare
        x = fluid.dygraph.to_variable(
            np.random.random((1, 784)).astype('float32'))
        self.assertTrue(
417 418
            np.array_equal(train_layer(x).numpy(),
                           new_layer(x).numpy()))
419

420
    def test_load_dygraph_no_path(self):
421 422
        model_path = os.path.join(self.temp_dir.name,
                                  "test_jit_save_load.no_path/model_path")
423 424 425
        with self.assertRaises(ValueError):
            model_dict, _ = fluid.dygraph.load_dygraph(model_path)

426
    def test_jit_load_no_path(self):
427 428
        path = os.path.join(self.temp_dir.name,
                            "test_jit_save_load.no_path/model_path")
429 430 431
        with self.assertRaises(ValueError):
            loaded_layer = paddle.jit.load(path)

432

433
class TestSaveLoadWithNestOut(unittest.TestCase):
434

435 436 437
    def setUp(self):
        # enable dygraph mode
        fluid.enable_dygraph()
438 439 440 441
        self.temp_dir = tempfile.TemporaryDirectory()

    def tearDown(self):
        self.temp_dir.cleanup()
442 443 444 445 446 447 448 449 450

    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')])

451
        model_path = os.path.join(self.temp_dir.name, "net_with_nest_out/model")
452 453 454 455 456 457 458 459 460 461
        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()))


462
class TestSaveLoadWithDictInput(unittest.TestCase):
463

464
    def test_dict_input(self):
465
        # NOTE: This net cannot be executed, it is just
466 467 468
        # a special case for exporting models in model validation
        # We DO NOT recommend this writing way of Layer
        net = LinearNetWithDictInput(8, 8)
469 470 471
        # net.forward.concrete_program.inputs:
        # (<__main__.LinearNetWithDictInput object at 0x7f2655298a98>,
        #  {'img': var img : fluid.VarType.LOD_TENSOR.shape(-1, 8).astype(VarType.FP32)},
472 473
        #  {'label': var label : fluid.VarType.LOD_TENSOR.shape(-1, 1).astype(VarType.INT64)})
        self.assertEqual(len(net.forward.concrete_program.inputs), 3)
474 475 476
        temp_dir = tempfile.TemporaryDirectory()
        path = os.path.join(temp_dir.name,
                            "test_jit_save_load_with_dict_input/model")
477
        # prune inputs
478 479 480 481 482 483 484 485
        paddle.jit.save(layer=net,
                        path=path,
                        input_spec=[{
                            'img':
                            InputSpec(shape=[None, 8],
                                      dtype='float32',
                                      name='img')
                        }])
486 487 488 489 490 491 492 493

        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)
494
        temp_dir.cleanup()
495 496


497
class TestSaveLoadWithDictInputNoPrune(unittest.TestCase):
498

499 500
    def test_dict_input(self):
        net = LinearNetWithDictInputNoPrune(8, 8)
501 502 503
        temp_dir = tempfile.TemporaryDirectory()
        path = os.path.join(
            temp_dir.name, "test_jit_save_load_with_dict_input_no_prune/model")
504
        # prune inputs
505 506 507 508 509 510 511 512 513 514 515 516
        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')
                        }])
517 518 519 520 521 522 523

        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)
524
        temp_dir.cleanup()
525 526


527
class TestSaveLoadWithInputSpec(unittest.TestCase):
528

529 530 531
    def setUp(self):
        # enable dygraph mode
        fluid.enable_dygraph()
532 533 534 535
        self.temp_dir = tempfile.TemporaryDirectory()

    def tearDown(self):
        self.temp_dir.cleanup()
536 537 538 539

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

543 544
        model_path = os.path.join(self.temp_dir.name,
                                  "input_spec.output_spec/model")
545 546 547 548 549 550 551
        # 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
552 553
        output_spec = net.forward.outputs[:1]
        paddle.jit.save(net, model_path, output_spec=output_spec)
554 555

        # 2. load to infer
556
        infer_layer = paddle.jit.load(model_path)
557 558 559 560 561 562 563
        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)

564 565
        model_path = os.path.join(self.temp_dir.name,
                                  "multi_inout.output_spec1/model")
566 567 568 569 570 571 572 573
        # 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
574 575
        output_spec = net.forward.outputs[:2]
        paddle.jit.save(net, model_path, output_spec=output_spec)
576 577

        # 3. load to infer
578
        infer_layer = paddle.jit.load(model_path)
579 580 581 582 583 584 585 586
        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
587 588
        model_path = os.path.join(self.temp_dir.name,
                                  "multi_inout.output_spec2/model")
589 590
        output_spec = net.forward.outputs[:1]
        paddle.jit.save(net, model_path, [input_x], output_spec=output_spec)
591
        # 2. load again
592
        infer_layer2 = paddle.jit.load(model_path)
593 594 595 596 597
        # 3. predict
        pred_xx = infer_layer2(x)

        # 4. assert pred_x == pred_xx
        self.assertTrue(np.allclose(pred_x.numpy(), pred_xx.numpy()))
598 599 600 601

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

602 603
        model_path = os.path.join(self.temp_dir.name,
                                  "multi_inout1.output_spec1/model")
604 605 606 607 608 609 610 611 612 613 614 615 616 617 618 619 620 621 622 623 624
        # 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
625 626
        model_path = os.path.join(self.temp_dir.name,
                                  "multi_inout1.output_spec2/model")
627 628 629 630 631 632 633 634 635
        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()))
636 637


638
class TestJitSaveLoadConfig(unittest.TestCase):
639

640 641 642 643
    def setUp(self):
        # enable dygraph mode
        fluid.enable_dygraph()
        # config seed
C
cnn 已提交
644
        paddle.seed(SEED)
L
Leo Chen 已提交
645
        paddle.framework.random._manual_program_seed(SEED)
646 647 648 649
        self.temp_dir = tempfile.TemporaryDirectory()

    def tearDown(self):
        self.temp_dir.cleanup()
650 651 652 653 654 655 656 657 658 659 660 661 662

    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()

663 664
        model_path = os.path.join(self.temp_dir.name,
                                  "save_load_config.output_spec")
665
        output_spec = [out]
666 667 668 669
        paddle.jit.save(layer=train_layer,
                        path=model_path,
                        input_spec=[x],
                        output_spec=output_spec)
670 671

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

679 680
    def test_save_no_support_config_error(self):
        layer = LinearNet(784, 1)
681
        path = os.path.join(self.temp_dir.name, "no_support_config_test")
682 683 684 685
        with self.assertRaises(ValueError):
            paddle.jit.save(layer=layer, path=path, model_filename="")

    def test_load_empty_model_filename_error(self):
686
        path = os.path.join(self.temp_dir.name, "error_model_filename_test")
687 688 689 690
        with self.assertRaises(ValueError):
            paddle.jit.load(path, model_filename="")

    def test_load_empty_params_filename_error(self):
691
        path = os.path.join(self.temp_dir.name, "error_params_filename_test")
692 693 694 695
        with self.assertRaises(ValueError):
            paddle.jit.load(path, params_filename="")

    def test_load_with_no_support_config(self):
696
        path = os.path.join(self.temp_dir.name, "no_support_config_test")
697 698 699
        with self.assertRaises(ValueError):
            paddle.jit.load(path, separate_params=True)

700

701
class TestJitMultipleLoading(unittest.TestCase):
702

703 704
    def setUp(self):
        self.linear_size = 4
705 706 707
        self.temp_dir = tempfile.TemporaryDirectory()
        self.model_path = os.path.join(self.temp_dir.name,
                                       "jit_multi_load/model")
708 709 710
        # enable dygraph mode
        fluid.enable_dygraph()
        # config seed
C
cnn 已提交
711
        paddle.seed(SEED)
L
Leo Chen 已提交
712
        paddle.framework.random._manual_program_seed(SEED)
713 714 715
        # train and save base model
        self.train_and_save_orig_model()

716 717 718
    def tearDown(self):
        self.temp_dir.cleanup()

719 720 721
    def train_and_save_orig_model(self):
        layer = LinearNet(self.linear_size, self.linear_size)
        example_inputs, layer, _ = train(layer, self.linear_size, 1)
722 723 724
        paddle.jit.save(layer=layer,
                        path=self.model_path,
                        input_spec=example_inputs)
725 726 727 728 729 730 731 732 733 734 735

    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)


736
class TestJitPruneModelAndLoad(unittest.TestCase):
737

738 739
    def setUp(self):
        self.linear_size = 4
740 741 742
        self.temp_dir = tempfile.TemporaryDirectory()
        self.model_path = os.path.join(self.temp_dir.name,
                                       "jit_prune_model_and_load/model")
743 744 745
        # enable dygraph mode
        fluid.enable_dygraph()
        # config seed
C
cnn 已提交
746
        paddle.seed(SEED)
L
Leo Chen 已提交
747
        paddle.framework.random._manual_program_seed(SEED)
748

749 750 751
    def tearDown(self):
        self.temp_dir.cleanup()

752 753 754 755 756 757 758 759 760 761 762 763
    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()

764
        output_spec = [hidden]
765 766 767 768
        paddle.jit.save(layer=train_layer,
                        path=self.model_path,
                        input_spec=[x],
                        output_spec=output_spec)
769 770 771 772 773 774 775

        return train_layer

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

776
        infer_layer = paddle.jit.load(self.model_path)
777 778 779 780

        x = fluid.dygraph.to_variable(
            np.random.random((4, 8)).astype('float32'))
        self.assertTrue(
781 782
            np.array_equal(train_layer(x)[0].numpy(),
                           infer_layer(x).numpy()))
783 784 785 786 787

    def test_load_var_not_in_extra_var_info(self):
        self.train_and_save()

        # chage extra var info
788
        var_info_path = self.model_path + INFER_PARAMS_INFO_SUFFIX
789 790 791 792 793 794 795
        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):
796
            paddle.jit.load(self.model_path)
797 798


799
class TestJitSaveMultiCases(unittest.TestCase):
800

801 802 803 804
    def setUp(self):
        # enable dygraph mode
        fluid.enable_dygraph()
        # config seed
C
cnn 已提交
805
        paddle.seed(SEED)
806
        paddle.framework.random._manual_program_seed(SEED)
807 808 809 810
        self.temp_dir = tempfile.TemporaryDirectory()

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

C
Chen Weihang 已提交
812 813 814 815 816
    def verify_inference_correctness(self,
                                     layer,
                                     model_path,
                                     with_label_and_loss=False,
                                     with_label=False):
817 818 819 820
        layer.eval()
        loaded_layer = paddle.jit.load(model_path)
        loaded_layer.eval()
        # inference & compare
Z
Zhou Wei 已提交
821
        x = paddle.to_tensor(np.random.random((1, 784)).astype('float32'))
C
Chen Weihang 已提交
822
        if with_label_and_loss:
Z
Zhou Wei 已提交
823
            y = paddle.to_tensor(np.random.random((1, 1)).astype('int64'))
824 825
            pred, _ = layer(x, y)
            pred = pred.numpy()
C
Chen Weihang 已提交
826 827 828 829
        elif with_label:
            y = paddle.to_tensor(np.random.random((1, 1)).astype('int64'))
            pred = layer(x, y)
            pred = pred.numpy()
830 831 832 833 834 835 836 837 838 839 840 841 842
        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)

843 844
        model_path = os.path.join(self.temp_dir.name,
                                  "test_no_prune_to_static_after_train/model")
845 846 847 848 849 850 851
        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)

852 853
        model_path = os.path.join(self.temp_dir.name,
                                  "test_no_prune_to_static_no_train/model")
854 855 856 857 858 859 860 861 862
        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)

863 864
        model_path = os.path.join(
            self.temp_dir.name, "test_no_prune_no_to_static_after_train/model")
865 866 867
        paddle.jit.save(
            layer,
            model_path,
868
            input_spec=[InputSpec(shape=[None, 784], dtype='float32')])
869 870 871 872 873 874 875 876

        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)

877 878 879
        model_path = os.path.join(
            self.temp_dir.name,
            "test_no_prune_no_to_static_after_train_with_examples/model")
880
        paddle.jit.save(layer=layer, path=model_path, input_spec=example_inputs)
881 882 883 884 885 886

        self.verify_inference_correctness(layer, model_path)

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

887 888
        model_path = os.path.join(self.temp_dir.name,
                                  "test_no_prune_no_to_static_no_train/model")
889 890 891
        paddle.jit.save(
            layer,
            model_path,
892
            input_spec=[InputSpec(shape=[None, 784], dtype='float32')])
893 894 895 896 897 898 899 900

        self.verify_inference_correctness(layer, model_path)

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

        out = train_with_label(layer)

901 902
        model_path = os.path.join(self.temp_dir.name,
                                  "test_prune_to_static_after_train/model")
903 904 905 906 907 908 909 910 911 912 913 914
        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)
915 916 917 918

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

919 920
        model_path = os.path.join(self.temp_dir.name,
                                  "test_prune_to_static_no_train/model")
921 922
        # TODO: no train, cannot get output_spec var here
        # now only can use index
923
        output_spec = layer.forward.outputs[:1]
924 925 926 927 928 929 930 931 932 933 934 935
        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)
C
Chen Weihang 已提交
936 937 938 939

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

940 941
        model_path = os.path.join(self.temp_dir.name,
                                  "test_prune_input_to_static_no_train/model")
942 943 944 945 946 947 948
        paddle.jit.save(layer,
                        model_path,
                        input_spec=[
                            InputSpec(shape=[None, 784],
                                      dtype='float32',
                                      name="image")
                        ])
C
Chen Weihang 已提交
949 950 951 952 953 954

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

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

955 956 957
        model_path = os.path.join(
            self.temp_dir.name,
            "test_prune_useless_input_to_static_no_train/model")
958 959 960 961 962 963 964
        paddle.jit.save(layer,
                        model_path,
                        input_spec=[
                            InputSpec(shape=[None, 784],
                                      dtype='float32',
                                      name="image")
                        ])
C
Chen Weihang 已提交
965 966

        self.verify_inference_correctness(layer, model_path, with_label=True)
967 968 969 970 971 972

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

        train(layer)

973 974
        model_path = os.path.join(
            self.temp_dir.name, "test_no_prune_input_spec_name_warning/model")
975 976 977
        paddle.jit.save(
            layer,
            model_path,
978 979 980 981 982 983 984 985
            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')
                        ])
986 987 988 989 990 991 992 993

        self.verify_inference_correctness(layer, model_path)

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

        train(layer)

994 995
        model_path = os.path.join(
            self.temp_dir.name, "test_not_prune_output_spec_name_warning/model")
Z
Zhou Wei 已提交
996
        out = paddle.to_tensor(np.random.random((1, 1)).astype('float'))
997
        paddle.jit.save(layer, model_path, output_spec=[out])
998 999 1000 1001 1002 1003

        self.verify_inference_correctness(layer, model_path)

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

1004 1005
        model_path = os.path.join(self.temp_dir.name,
                                  "test_prune_input_spec_name_error/model")
1006 1007 1008 1009
        with self.assertRaises(ValueError):
            paddle.jit.save(
                layer,
                model_path,
1010
                input_spec=[InputSpec(shape=[None, 784], dtype='float32')])
1011
        with self.assertRaises(ValueError):
1012 1013 1014 1015 1016 1017 1018
            paddle.jit.save(layer,
                            model_path,
                            input_spec=[
                                InputSpec(shape=[None, 784],
                                          dtype='float32',
                                          name='feed_input')
                            ])
1019 1020 1021 1022 1023 1024

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

        train_with_label(layer)

1025 1026
        model_path = os.path.join(self.temp_dir.name,
                                  "test_prune_to_static_after_train/model")
Z
Zhou Wei 已提交
1027
        out = paddle.to_tensor(np.random.random((1, 1)).astype('float'))
1028
        with self.assertRaises(ValueError):
1029 1030 1031 1032 1033 1034 1035 1036
            paddle.jit.save(layer,
                            model_path,
                            input_spec=[
                                InputSpec(shape=[None, 784],
                                          dtype='float32',
                                          name="image")
                            ],
                            output_spec=[out])
1037 1038


1039
class TestJitSaveLoadEmptyLayer(unittest.TestCase):
1040

1041
    def setUp(self):
1042 1043 1044
        self.temp_dir = tempfile.TemporaryDirectory()
        self.model_path = os.path.join(self.temp_dir.name,
                                       "jit_save_load_empty_layer/model")
1045 1046 1047
        # enable dygraph mode
        paddle.disable_static()

1048 1049 1050
    def tearDown(self):
        self.temp_dir.cleanup()

1051 1052
    def test_save_load_empty_layer(self):
        layer = EmptyLayer()
Z
Zhou Wei 已提交
1053
        x = paddle.to_tensor(np.random.random((10)).astype('float32'))
1054 1055 1056 1057 1058 1059 1060 1061
        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):
1062

1063
    def setUp(self):
1064 1065 1066
        self.temp_dir = tempfile.TemporaryDirectory()
        self.model_path = os.path.join(self.temp_dir.name,
                                       "jit_save_load_no_param_layer/model")
1067 1068 1069
        # enable dygraph mode
        paddle.disable_static()

1070 1071 1072
    def tearDown(self):
        self.temp_dir.cleanup()

1073 1074
    def test_save_load_no_param_layer(self):
        layer = NoParamLayer()
Z
Zhou Wei 已提交
1075 1076
        x = paddle.to_tensor(np.random.random((5)).astype('float32'))
        y = paddle.to_tensor(np.random.random((5)).astype('float32'))
1077 1078 1079 1080 1081 1082 1083
        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))


1084
class TestJitSaveLoadMultiMethods(unittest.TestCase):
1085

1086 1087 1088
    def setUp(self):
        # enable dygraph mode
        paddle.disable_static()
1089 1090 1091 1092
        self.temp_dir = tempfile.TemporaryDirectory()

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

    def test_jit_save_load_inference(self):
1095 1096
        model_path_inference = os.path.join(
            self.temp_dir.name, "jit_save_load_multi_methods/model")
1097 1098 1099 1100 1101 1102 1103 1104 1105 1106 1107
        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(
1108 1109
                float((result -
                       getattr(load_net, func, None)(inps)).abs().max()) < 1e-5)
1110 1111

    def test_jit_save_load_multi_methods_inputspec(self):
1112 1113
        model_path = os.path.join(self.temp_dir.name,
                                  'jit_save_load_multi_methods/model')
1114 1115
        layer = LinearNetWithMultiStaticFunc(784, 1)
        with self.assertRaises(ValueError):
1116 1117 1118
            paddle.jit.save(layer,
                            model_path,
                            input_spec=[InputSpec(shape=[None, 784])])
1119

1120
    def test_parse_name(self):
1121 1122
        model_path_inference = os.path.join(self.temp_dir.name,
                                            "jit_save_load_parse_name/model")
1123 1124 1125 1126 1127 1128 1129 1130 1131 1132
        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'))

1133

W
WeiXin 已提交
1134
class LayerSaved(paddle.nn.Layer):
1135

W
WeiXin 已提交
1136 1137 1138 1139 1140 1141 1142 1143 1144 1145 1146 1147 1148
    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
1149
        if paddle.shape(x)[0] == 1:
W
WeiXin 已提交
1150 1151 1152 1153 1154 1155 1156
            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):
1157

W
WeiXin 已提交
1158 1159 1160 1161 1162 1163 1164 1165 1166 1167 1168 1169 1170 1171 1172 1173 1174 1175
    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
1176
        if paddle.shape(x)[0] == 1:
W
WeiXin 已提交
1177 1178 1179 1180 1181 1182 1183 1184 1185 1186 1187 1188 1189
            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


1190
class TestJitSaveLoadSaveWithoutRunning(unittest.TestCase):
1191

1192 1193 1194
    def setUp(self):
        # enable dygraph mode
        paddle.disable_static()
1195 1196 1197 1198
        self.temp_dir = tempfile.TemporaryDirectory()

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

    def test_save_load_finetune_load(self):
1201 1202
        model_path = os.path.join(
            self.temp_dir.name, "test_jit_save_load_save_without_running/model")
1203 1204 1205 1206 1207 1208 1209
        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
1210 1211 1212 1213 1214 1215
        paddle.jit.save(layer_save,
                        model_path,
                        input_spec=[
                            paddle.static.InputSpec(shape=[None, IMAGE_SIZE],
                                                    dtype='float32')
                        ])
1216 1217 1218 1219 1220
        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)
1221 1222 1223 1224 1225 1226
            paddle.jit.save(layer_load,
                            model_path,
                            input_spec=[
                                paddle.static.InputSpec(
                                    shape=[None, IMAGE_SIZE], dtype='float32')
                            ])
1227 1228 1229 1230 1231 1232 1233 1234 1235
        #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)


W
WeiXin 已提交
1236
class TestJitSaveLoadFinetuneLoad(unittest.TestCase):
1237

W
WeiXin 已提交
1238 1239 1240
    def setUp(self):
        # enable dygraph mode
        paddle.disable_static()
1241 1242 1243 1244
        self.temp_dir = tempfile.TemporaryDirectory()

    def tearDown(self):
        self.temp_dir.cleanup()
W
WeiXin 已提交
1245 1246

    def test_save_load_finetune_load(self):
1247 1248
        model_path = os.path.join(self.temp_dir.name,
                                  "test_jit_save_load_finetune_load/model")
W
WeiXin 已提交
1249 1250 1251 1252 1253 1254 1255 1256 1257 1258 1259 1260 1261 1262 1263 1264 1265 1266 1267 1268 1269 1270 1271 1272 1273 1274 1275
        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)


1276 1277 1278 1279
# 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,
1280 1281
# TestJitSaveLoadFunctionCase2, TestJitSaveLoadFunctionCase3.
class TestJitSaveLoadFunctionCase1(unittest.TestCase):
1282

1283 1284
    def setUp(self):
        paddle.disable_static()
1285 1286 1287 1288
        self.temp_dir = tempfile.TemporaryDirectory()

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

    def test_jit_save_load_static_function(self):
1291

1292 1293 1294 1295
        @paddle.jit.to_static
        def fun(inputs):
            return paddle.tanh(inputs)

1296 1297
        path = os.path.join(self.temp_dir.name,
                            'test_jit_save_load_function_1/func')
1298 1299 1300 1301 1302 1303 1304 1305 1306
        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)

1307 1308

class TestJitSaveLoadFunctionCase2(unittest.TestCase):
1309

1310 1311
    def setUp(self):
        paddle.disable_static()
1312 1313 1314 1315
        self.temp_dir = tempfile.TemporaryDirectory()

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

1317
    def test_jit_save_load_function_input_spec(self):
1318

1319
        @paddle.jit.to_static(input_spec=[
1320
            InputSpec(shape=[None, 6], dtype='float32', name='x'),
1321 1322 1323 1324
        ])
        def fun(inputs):
            return paddle.nn.functional.relu(inputs)

1325 1326
        path = os.path.join(self.temp_dir.name,
                            'test_jit_save_load_function_2/func')
1327 1328 1329 1330 1331 1332 1333 1334
        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)

1335 1336

class TestJitSaveLoadFunctionCase3(unittest.TestCase):
1337

1338 1339
    def setUp(self):
        paddle.disable_static()
1340 1341 1342 1343
        self.temp_dir = tempfile.TemporaryDirectory()

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

1345
    def test_jit_save_load_function_function(self):
1346

1347 1348 1349
        def fun(inputs):
            return paddle.tanh(inputs)

1350 1351
        path = os.path.join(self.temp_dir.name,
                            'test_jit_save_load_function_3/func')
1352 1353 1354
        inps = paddle.rand([3, 6])
        origin = fun(inps)

1355 1356 1357 1358 1359 1360 1361
        paddle.jit.save(fun,
                        path,
                        input_spec=[
                            InputSpec(shape=[None, 6],
                                      dtype='float32',
                                      name='x'),
                        ])
1362 1363 1364 1365 1366 1367
        load_func = paddle.jit.load(path)

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


1368
class TestJitSaveLoadFunctionWithParamCase1(unittest.TestCase):
1369

1370 1371
    def setUp(self):
        paddle.disable_static()
1372 1373 1374 1375
        self.temp_dir = tempfile.TemporaryDirectory()

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

    def test_jit_save_load_function(self):
1378

1379
        class LinearNet(paddle.nn.Layer):
1380

1381 1382 1383 1384 1385 1386 1387 1388 1389 1390 1391 1392 1393 1394 1395
            def __init__(self):
                super(LinearNet, self).__init__()
                self._linear = paddle.nn.Linear(5, 6)

            def forward(self, x):
                return paddle.tanh(x)

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

        layer = LinearNet()

        inps = paddle.rand([3, 5])
        origin = layer.anothor_forward(inps)

1396 1397
        func = paddle.jit.to_static(layer.anothor_forward,
                                    [paddle.static.InputSpec(shape=[-1, 5])])
1398 1399 1400
        path = os.path.join(
            self.temp_dir.name,
            'test_jit_save_load_function_with_params_case1/func')
1401 1402 1403 1404 1405 1406 1407 1408
        paddle.jit.save(func, path)
        load_func = paddle.jit.load(path)

        load_result = load_func(inps)
        self.assertTrue(np.array_equal(load_result.numpy(), origin.numpy()))


class TestJitSaveLoadFunctionWithParamCase2(unittest.TestCase):
1409

1410 1411
    def setUp(self):
        paddle.disable_static()
1412 1413 1414 1415
        self.temp_dir = tempfile.TemporaryDirectory()

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

    def test_jit_save_load_function(self):
1418

1419
        class LinearNet(paddle.nn.Layer):
1420

1421 1422 1423 1424 1425 1426 1427 1428 1429 1430 1431 1432 1433 1434 1435
            def __init__(self):
                super(LinearNet, self).__init__()
                self._linear = paddle.nn.Linear(5, 6)

            def forward(self, x):
                return paddle.tanh(x)

            @paddle.jit.to_static(input_spec=[InputSpec(shape=[-1, 5])])
            def anothor_forward(self, x):
                return self._linear(x)

        layer = LinearNet()

        inps = paddle.rand([3, 5])

1436 1437 1438
        path = os.path.join(
            self.temp_dir.name,
            'test_jit_save_load_function_with_params_case2/func')
1439 1440 1441 1442 1443 1444 1445 1446 1447 1448 1449
        paddle.jit.save(layer.anothor_forward, path)
        origin_result = layer.anothor_forward(inps)
        load_func = paddle.jit.load(path)

        load_result = load_func(inps)

        self.assertTrue(
            np.array_equal(origin_result.numpy(), load_result.numpy()))


class TestJitSaveLoadFunctionWithParamCase3(unittest.TestCase):
1450

1451 1452
    def setUp(self):
        paddle.disable_static()
1453 1454 1455 1456
        self.temp_dir = tempfile.TemporaryDirectory()

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

    def test_jit_save_load_function(self):
1459

1460
        class LinearNet(paddle.nn.Layer):
1461

1462 1463 1464 1465 1466 1467 1468 1469 1470 1471 1472 1473 1474 1475 1476 1477
            def __init__(self):
                super(LinearNet, self).__init__()
                self._linear = paddle.nn.Linear(5, 6)

            def forward(self, x):
                return paddle.tanh(x)

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

        layer = LinearNet()

        inps = paddle.rand([3, 5])
        origin = layer.anothor_forward(inps)

1478 1479 1480
        path = os.path.join(
            self.temp_dir.name,
            'test_jit_save_load_function_with_params_case3/func')
1481 1482 1483 1484 1485 1486 1487
        paddle.jit.save(layer.anothor_forward, path)
        load_func = paddle.jit.load(path)

        load_result = load_func(inps)
        self.assertTrue(np.array_equal(load_result.numpy(), origin.numpy()))


1488
class TestJitSaveLoadDataParallel(unittest.TestCase):
1489

1490 1491 1492 1493 1494 1495
    def setUp(self):
        self.temp_dir = tempfile.TemporaryDirectory()

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

1496 1497 1498 1499 1500 1501 1502 1503 1504 1505 1506 1507 1508 1509 1510 1511
    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)
1512 1513
        path = os.path.join(self.temp_dir.name,
                            "jit_save_data_parallel_with_inputspec/model")
1514 1515 1516
        paddle.jit.save(layer=layer,
                        path=path,
                        input_spec=[InputSpec(shape=[None, 784])])
1517 1518 1519 1520 1521 1522 1523

        self.verify_inference_correctness(layer, path)

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

1524 1525
        path = os.path.join(self.temp_dir.name,
                            "jit_save_data_parallel_with_to_static/model")
1526 1527 1528 1529 1530
        paddle.jit.save(layer, path)

        self.verify_inference_correctness(layer, path)


1531 1532 1533 1534 1535 1536
class InputSepcLayer(paddle.nn.Layer):
    '''
    A layer with InputSpec to test InputSpec compatibility
    '''

    @paddle.jit.to_static(input_spec=[
1537 1538
        InputSpec(shape=[None, 8], dtype='float32', name='x'),
        InputSpec(shape=[None, 1], dtype='float64', name='y')
1539 1540 1541 1542 1543 1544
    ])
    def forward(self, x, y):
        return x, y


class TestInputSpecCompatibility(unittest.TestCase):
1545

1546 1547 1548 1549 1550 1551
    def setUp(self):
        self.temp_dir = tempfile.TemporaryDirectory()

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

1552 1553 1554 1555 1556 1557 1558 1559 1560 1561 1562 1563
    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()
1564 1565
        save_dir = os.path.join(self.temp_dir.name,
                                "jit_save_compatible_input_spec")
1566 1567 1568 1569 1570 1571 1572
        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)

1573 1574 1575 1576 1577 1578 1579 1580 1581 1582
        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')
                        ])
1583 1584 1585 1586
        same_input_spec_layer = paddle.jit.load(path)
        self._assert_input_spec_layer_return(layer, same_input_spec_layer)
        shutil.rmtree(save_dir)

1587 1588 1589 1590 1591 1592
        paddle.jit.save(layer=layer,
                        path=path,
                        input_spec=[
                            InputSpec(shape=[8, 8], dtype='float32'),
                            InputSpec(shape=[8, -1], dtype='float64')
                        ])
1593 1594 1595 1596 1597 1598
        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()
1599 1600
        save_dir = os.path.join(self.temp_dir.name,
                                "jit_save_compatible_input_spec")
1601 1602 1603 1604
        path = save_dir + "/model"

        with self.assertRaises(ValueError):
            # type mismatch
1605 1606 1607 1608 1609 1610
            paddle.jit.save(layer=layer,
                            path=path,
                            input_spec=[
                                InputSpec(shape=[None, 8], dtype='float64'),
                                InputSpec(shape=[None, 1], dtype='float64')
                            ])
1611 1612 1613

        with self.assertRaises(ValueError):
            # shape len mismatch
1614 1615 1616 1617 1618 1619
            paddle.jit.save(layer=layer,
                            path=path,
                            input_spec=[
                                InputSpec(shape=[None, 8, 1], dtype='float32'),
                                InputSpec(shape=[None, 1], dtype='float64')
                            ])
1620 1621 1622

        with self.assertRaises(ValueError):
            # shape mismatch
1623 1624 1625 1626 1627 1628
            paddle.jit.save(layer=layer,
                            path=path,
                            input_spec=[
                                InputSpec(shape=[None, 8], dtype='float32'),
                                InputSpec(shape=[None, 2], dtype='float64')
                            ])
1629 1630 1631 1632
        if os.path.exists(save_dir):
            shutil.rmtree(save_dir)


1633
if __name__ == '__main__':
1634 1635
    with fluid.framework._test_eager_guard():
        unittest.main()