Unset ReserveSpace of batch_norm for inference program.

[Xreki]

PR types

Bug fixes

PR changes

APIs

Describe

• batch_norm的ReserveSpace输出只在训练时需要使用，用以保存一些前向计算的临时结果给反向使用。当前save_inference_model因保存了ReserveSpace，导致Lite模型batch_norm相关的pass匹配失败。因此在save_inference_model接口保存预测模型时，将ReserveSpace输出移除。

• 因有将预测模型加载进来，使用动态图重新训练的需求。所以在将预测program扩展支持训练时，需要将ReserveSpace输出加回来。

• 测试代码

import paddle
import paddle.nn as nn
import numpy as np


class Model(nn.Layer):
    def __init__(self, num_channels, use_global_stats=None):
        super(Model, self).__init__()
        self.conv2d_0 = nn.Conv2D(
            in_channels=num_channels,
            out_channels=6,
            kernel_size=3,
            bias_attr=False,
            data_format='NCHW')
        self.batch_norm_0 = nn.BatchNorm2D(
            num_features=6, data_format='NCHW', use_global_stats=use_global_stats)

        self.conv2d_1 = nn.Conv2D(
            in_channels=6,
            out_channels=8,
            kernel_size=1,
            bias_attr=False,
            data_format='NCHW')
        self.batch_norm_1 = nn.BatchNorm2D(
            num_features=8, data_format='NCHW', use_global_stats=use_global_stats)

        self.in_features = 5408
        self.linear = nn.Linear(in_features=self.in_features, out_features=10)

    def forward(self, input):
        conv_out_0 = self.conv2d_0(input)
        bn_out_0 = self.batch_norm_0(conv_out_0)
        conv_out_1 = self.conv2d_1(bn_out_0)
        bn_out_1 = self.batch_norm_1(conv_out_1)
        bn_out_1 = paddle.reshape(bn_out_1, shape=[-1, self.in_features])
        linear_out = self.linear(bn_out_1)
        out = nn.functional.softmax(linear_out)
        return out


def build_and_save_model(training, path_prefix):
    paddle.enable_static()

    num_channels = 1
    model = Model(num_channels)
    if training:
        model.train()
    else:
        model.eval()

    input = paddle.static.data(name='input', shape=[None, num_channels, 28, 28], dtype='float32')
    predict = model(input)
    loss = paddle.mean(predict)

    optimizer = paddle.optimizer.Momentum()
    optimizer.minimize(loss)

    exe = paddle.static.Executor(paddle.CPUPlace())
    exe.run(paddle.static.default_startup_program())

    paddle.static.save_inference_model(path_prefix, [input], [predict], exe)


def load_and_train(mode="dynamic"):
    path_prefix = "prune_bn"
    build_and_save_model(True, path_prefix)

    if mode == "dynamic":
        paddle.disable_static()

        model = paddle.jit.load(path_prefix)
        model.train()

        data = np.random.random([1, 1, 28, 28]).astype("float32")
        input = paddle.to_tensor(data)
        predict = model(input)

        loss = paddle.mean(predict)
        loss.backward()
        adam = paddle.optimizer.Adam(learning_rate=0.001, parameters=model.parameters())
        adam.step()
        adam.clear_grad()
    else:
        exe = paddle.static.Executor(paddle.CPUPlace())
        [inference_program, feed_target_names, fetch_targets] = paddle.static.load_inference_model(path_prefix, exe)
        print(inference_program)


if __name__ == '__main__':
    load_and_train(mode="static")

inference_program模型结构打印如下：

模型可视化如下：

[paddle-bot-old]

Thanks for your contribution!
Please wait for the result of CI firstly. See Paddle CI Manual for details.

[chenwhql]

LGTM

[chenwhql]

小问题，后面是不是还有个3的注释，要不要也顺便更新下

[Xreki]

看漏了，隔了几行代码才有3，我在后续别的PR里面更新吧。

[hbwx24]

LGTM