test_quantization_pass.py 27.9 KB
Newer Older
W
WangZhen 已提交
1 2 3 4 5 6 7 8 9 10 11 12 13 14
#   copyright (c) 2018 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.

Z
Zhen Wang 已提交
15
import os
W
WangZhen 已提交
16 17 18 19 20
import unittest
import random
import numpy as np
import paddle.fluid as fluid
import six
W
WangZhen 已提交
21
import paddle
22
from paddle.fluid.framework import IrGraph
23
from paddle.fluid.contrib.slim.quantization import QuantizationTransformPass
W
WangZhen 已提交
24
from paddle.fluid.contrib.slim.quantization import QuantizationFreezePass
25 26
from paddle.fluid.contrib.slim.quantization import ConvertToInt8Pass
from paddle.fluid.contrib.slim.quantization import TransformForMobilePass
27
from paddle.fluid.contrib.slim.quantization import AddQuantDequantPass
W
WangZhen 已提交
28 29
from paddle.fluid import core

P
pangyoki 已提交
30 31
paddle.enable_static()

Z
Zhen Wang 已提交
32 33 34
os.environ["CUDA_VISIBLE_DEVICES"] = "0"
os.environ["CPU_NUM"] = "1"

W
WangZhen 已提交
35 36 37 38 39 40 41 42 43 44 45 46

def linear_fc(num):
    data = fluid.layers.data(name='image', shape=[1, 32, 32], dtype='float32')
    label = fluid.layers.data(name='label', shape=[1], dtype='int64')
    hidden = data
    for _ in six.moves.xrange(num):
        hidden = fluid.layers.fc(hidden, size=128, act='relu')
    loss = fluid.layers.cross_entropy(input=hidden, label=label)
    loss = fluid.layers.mean(loss)
    return loss


47
def residual_block(num, quant_skip_pattern=None):
W
WangZhen 已提交
48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64
    def conv_bn_layer(input,
                      ch_out,
                      filter_size,
                      stride,
                      padding,
                      act='relu',
                      bias_attr=False):
        tmp = fluid.layers.conv2d(
            input=input,
            filter_size=filter_size,
            num_filters=ch_out,
            stride=stride,
            padding=padding,
            act=None,
            bias_attr=bias_attr)
        return fluid.layers.batch_norm(input=tmp, act=act)

65 66 67 68 69 70 71
    data = fluid.layers.data(
        name='image',
        shape=[1, 1, 32, 32],
        dtype='float32',
        append_batch_size=False)
    label = fluid.layers.data(
        name='label', shape=[1, 1], dtype='int64', append_batch_size=False)
W
WangZhen 已提交
72 73 74 75 76
    hidden = data
    for _ in six.moves.xrange(num):
        conv = conv_bn_layer(hidden, 16, 3, 1, 1, act=None, bias_attr=True)
        short = conv_bn_layer(hidden, 16, 1, 1, 0, act=None)
        hidden = fluid.layers.elementwise_add(x=conv, y=short, act='relu')
77 78 79
    matmul_weight = fluid.layers.create_parameter(
        shape=[1, 16, 32, 32], dtype='float32')
    hidden = fluid.layers.matmul(hidden, matmul_weight, True, True)
80 81 82 83 84 85 86
    if quant_skip_pattern:
        with fluid.name_scope(quant_skip_pattern):
            pool = fluid.layers.pool2d(
                input=hidden, pool_size=2, pool_type='avg', pool_stride=2)
    else:
        pool = fluid.layers.pool2d(
            input=hidden, pool_size=2, pool_type='avg', pool_stride=2)
87
    fc = fluid.layers.fc(input=pool, size=10)
W
WangZhen 已提交
88 89 90 91 92
    loss = fluid.layers.cross_entropy(input=fc, label=label)
    loss = fluid.layers.mean(loss)
    return loss


93
def conv_net(img, label, quant_skip_pattern):
W
WangZhen 已提交
94 95 96 97 98 99
    conv_pool_1 = fluid.nets.simple_img_conv_pool(
        input=img,
        filter_size=5,
        num_filters=20,
        pool_size=2,
        pool_stride=2,
100
        pool_type='max',
W
WangZhen 已提交
101 102 103 104 105 106 107 108
        act="relu")
    conv_pool_1 = fluid.layers.batch_norm(conv_pool_1)
    conv_pool_2 = fluid.nets.simple_img_conv_pool(
        input=conv_pool_1,
        filter_size=5,
        num_filters=50,
        pool_size=2,
        pool_stride=2,
109
        pool_type='avg',
W
WangZhen 已提交
110
        act="relu")
111 112 113
    hidden = fluid.layers.fc(input=conv_pool_2, size=100, act='relu')
    with fluid.name_scope(quant_skip_pattern):
        prediction = fluid.layers.fc(input=hidden, size=10, act='softmax')
W
WangZhen 已提交
114 115 116 117 118
    loss = fluid.layers.cross_entropy(input=prediction, label=label)
    avg_loss = fluid.layers.mean(loss)
    return avg_loss


119
class TestQuantizationTransformPass(unittest.TestCase):
W
WangZhen 已提交
120 121 122 123 124 125
    def setUp(self):
        self.quantizable_op_and_inputs = {
            'conv2d': ['Input', 'Filter'],
            'depthwise_conv2d': ['Input', 'Filter'],
            'mul': ['X', 'Y']
        }
126
        self.quantizable_grad_op_inputs = {
W
WangZhen 已提交
127 128 129 130 131
            'conv2d_grad': ['Input', 'Filter'],
            'depthwise_conv2d_grad': ['Input', 'Filter'],
            'mul_grad': ['X', 'Y']
        }

132
    def check_program(self, program):
133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151
        quantized_ops = set()
        for block in program.blocks:
            for op in block.ops:
                # check forward
                if op.type in self.quantizable_op_and_inputs:
                    for arg_name in op.input_arg_names:
                        self.assertTrue(
                            arg_name.endswith('.quantized.dequantized'))
                        quantized_ops.add(arg_name)

            for op in block.ops:
                # check backward
                if op.type in self.quantizable_grad_op_inputs:
                    for pname in self.quantizable_grad_op_inputs[op.type]:
                        arg_name = op.input(pname)[0]
                        self.assertTrue(
                            arg_name.endswith('.quantized.dequantized'))
                        self.assertTrue(arg_name in quantized_ops)

152 153 154 155
    def linear_fc_quant(self,
                        activation_quant_type,
                        weight_quantize_type,
                        for_ci=True):
W
WangZhen 已提交
156 157 158 159 160 161
        main = fluid.Program()
        startup = fluid.Program()
        with fluid.program_guard(main, startup):
            loss = linear_fc(3)
            opt = fluid.optimizer.Adam(learning_rate=0.001)
            opt.minimize(loss)
162
        place = fluid.CPUPlace()
163
        graph = IrGraph(core.Graph(main.desc), for_test=False)
164 165
        transform_pass = QuantizationTransformPass(
            scope=fluid.global_scope(),
166
            place=place,
167 168
            activation_quantize_type=activation_quant_type,
            weight_quantize_type=weight_quantize_type)
169
        transform_pass.apply(graph)
Z
Zhen Wang 已提交
170
        if not for_ci:
Z
Zhen Wang 已提交
171 172 173 174
            marked_nodes = set()
            for op in graph.all_op_nodes():
                if op.name().find('quantize') > -1:
                    marked_nodes.add(op)
175 176
            graph.draw('.', 'quantize_fc_' + activation_quant_type,
                       marked_nodes)
177
        program = graph.to_program()
178
        self.check_program(program)
179
        val_graph = IrGraph(core.Graph(program.desc), for_test=False)
Z
Zhen Wang 已提交
180
        if not for_ci:
Z
Zhen Wang 已提交
181 182 183 184
            val_marked_nodes = set()
            for op in val_graph.all_op_nodes():
                if op.name().find('quantize') > -1:
                    val_marked_nodes.add(op)
185 186
            val_graph.draw('.', 'val_fc_' + activation_quant_type,
                           val_marked_nodes)
W
WangZhen 已提交
187

188
    def test_linear_fc_quant_abs_max(self):
189
        self.linear_fc_quant('abs_max', 'abs_max', for_ci=True)
W
WangZhen 已提交
190

191
    def test_linear_fc_quant_range_abs_max(self):
192
        self.linear_fc_quant('range_abs_max', 'abs_max', for_ci=True)
W
WangZhen 已提交
193

194
    def test_linear_fc_quant_moving_average_abs_max(self):
195 196
        self.linear_fc_quant(
            'moving_average_abs_max', 'channel_wise_abs_max', for_ci=True)
197

198 199 200
    def residual_block_quant(self,
                             activation_quant_type,
                             weight_quantize_type,
201
                             quantizable_op_type,
202
                             for_ci=True):
W
WangZhen 已提交
203 204 205 206 207 208
        main = fluid.Program()
        startup = fluid.Program()
        with fluid.program_guard(main, startup):
            loss = residual_block(2)
            opt = fluid.optimizer.Adam(learning_rate=0.001)
            opt.minimize(loss)
209
        place = fluid.CPUPlace()
210
        graph = IrGraph(core.Graph(main.desc), for_test=False)
211 212
        transform_pass = QuantizationTransformPass(
            scope=fluid.global_scope(),
213
            place=place,
214
            activation_quantize_type=activation_quant_type,
215 216
            weight_quantize_type=weight_quantize_type,
            quantizable_op_type=quantizable_op_type)
217
        transform_pass.apply(graph)
Z
Zhen Wang 已提交
218
        if not for_ci:
Z
Zhen Wang 已提交
219 220 221 222
            marked_nodes = set()
            for op in graph.all_op_nodes():
                if op.name().find('quantize') > -1:
                    marked_nodes.add(op)
223 224
            graph.draw('.', 'quantize_residual_' + activation_quant_type,
                       marked_nodes)
225
        program = graph.to_program()
226
        self.check_program(program)
227
        val_graph = IrGraph(core.Graph(program.desc), for_test=False)
Z
Zhen Wang 已提交
228
        if not for_ci:
Z
Zhen Wang 已提交
229 230 231 232
            val_marked_nodes = set()
            for op in val_graph.all_op_nodes():
                if op.name().find('quantize') > -1:
                    val_marked_nodes.add(op)
233 234
            val_graph.draw('.', 'val_residual_' + activation_quant_type,
                           val_marked_nodes)
W
WangZhen 已提交
235

236
    def test_residual_block_abs_max(self):
237 238 239
        quantizable_op_type = ['conv2d', 'depthwise_conv2d', 'mul', 'matmul']
        self.residual_block_quant(
            'abs_max', 'abs_max', quantizable_op_type, for_ci=True)
W
WangZhen 已提交
240

241
    def test_residual_block_range_abs_max(self):
242 243 244
        quantizable_op_type = ['conv2d', 'depthwise_conv2d', 'mul', 'matmul']
        self.residual_block_quant(
            'range_abs_max', 'abs_max', quantizable_op_type, for_ci=True)
W
WangZhen 已提交
245

246
    def test_residual_block_moving_average_abs_max(self):
247
        quantizable_op_type = ['conv2d', 'depthwise_conv2d', 'mul', 'matmul']
248
        self.residual_block_quant(
249 250 251 252
            'moving_average_abs_max',
            'channel_wise_abs_max',
            quantizable_op_type,
            for_ci=True)
253

W
WangZhen 已提交
254

W
WangZhen 已提交
255
class TestQuantizationFreezePass(unittest.TestCase):
256 257 258 259 260
    def freeze_graph(self,
                     use_cuda,
                     seed,
                     activation_quant_type,
                     weight_quant_type='abs_max',
261 262
                     for_ci=True,
                     quant_skip_pattern='skip_quant'):
W
WangZhen 已提交
263 264 265 266 267 268 269 270 271
        def build_program(main, startup, is_test):
            main.random_seed = seed
            startup.random_seed = seed
            with fluid.unique_name.guard():
                with fluid.program_guard(main, startup):
                    img = fluid.layers.data(
                        name='image', shape=[1, 28, 28], dtype='float32')
                    label = fluid.layers.data(
                        name='label', shape=[1], dtype='int64')
272
                    loss = conv_net(img, label, quant_skip_pattern)
W
WangZhen 已提交
273 274 275 276 277 278 279 280 281 282 283 284 285 286 287
                    if not is_test:
                        opt = fluid.optimizer.Adam(learning_rate=0.001)
                        opt.minimize(loss)
            return [img, label], loss

        random.seed(0)
        np.random.seed(0)

        main = fluid.Program()
        startup = fluid.Program()
        test_program = fluid.Program()
        feeds, loss = build_program(main, startup, False)
        build_program(test_program, startup, True)
        test_program = test_program.clone(for_test=True)
        main_graph = IrGraph(core.Graph(main.desc), for_test=False)
W
WangZhen 已提交
288
        test_graph = IrGraph(core.Graph(test_program.desc), for_test=True)
W
WangZhen 已提交
289 290 291

        place = fluid.CUDAPlace(0) if use_cuda else fluid.CPUPlace()
        exe = fluid.Executor(place)
W
WangZhen 已提交
292 293 294
        scope = fluid.Scope()
        with fluid.scope_guard(scope):
            exe.run(startup)
W
WangZhen 已提交
295
        transform_pass = QuantizationTransformPass(
296 297
            scope=scope,
            place=place,
298
            activation_quantize_type=activation_quant_type,
299 300
            weight_quantize_type=weight_quant_type,
            skip_pattern=quant_skip_pattern)
W
WangZhen 已提交
301 302
        transform_pass.apply(main_graph)
        transform_pass.apply(test_graph)
303
        dev_name = '_gpu_' if use_cuda else '_cpu_'
Z
Zhen Wang 已提交
304
        if not for_ci:
Z
Zhen Wang 已提交
305 306 307 308
            marked_nodes = set()
            for op in main_graph.all_op_nodes():
                if op.name().find('quantize') > -1:
                    marked_nodes.add(op)
309 310
            main_graph.draw('.', 'main' + dev_name + activation_quant_type + '_'
                            + weight_quant_type, marked_nodes)
Z
Zhen Wang 已提交
311 312 313 314
            marked_nodes = set()
            for op in test_graph.all_op_nodes():
                if op.name().find('quantize') > -1:
                    marked_nodes.add(op)
315 316
            test_graph.draw('.', 'test' + dev_name + activation_quant_type + '_'
                            + weight_quant_type, marked_nodes)
W
WangZhen 已提交
317

Z
Zhen Wang 已提交
318 319 320
        build_strategy = fluid.BuildStrategy()
        build_strategy.memory_optimize = False
        build_strategy.enable_inplace = False
321
        build_strategy.fuse_all_reduce_ops = False
Z
Zhen Wang 已提交
322 323
        binary = fluid.CompiledProgram(main_graph.graph).with_data_parallel(
            loss_name=loss.name, build_strategy=build_strategy)
324
        quantized_test_program = test_graph.to_program()
325
        iters = 5
326
        batch_size = 8
W
WangZhen 已提交
327 328 329 330 331 332 333 334

        train_reader = paddle.batch(
            paddle.reader.shuffle(
                paddle.dataset.mnist.train(), buf_size=500),
            batch_size=batch_size)
        test_reader = paddle.batch(
            paddle.dataset.mnist.test(), batch_size=batch_size)
        feeder = fluid.DataFeeder(feed_list=feeds, place=place)
W
WangZhen 已提交
335
        with fluid.scope_guard(scope):
W
WangZhen 已提交
336 337
            for _ in range(iters):
                data = next(train_reader())
Z
Zhen Wang 已提交
338
                loss_v = exe.run(binary,
339 340
                                 feed=feeder.feed(data),
                                 fetch_list=[loss])
Z
Zhen Wang 已提交
341
                if not for_ci:
342 343 344
                    print('{}: {}'.format('loss' + dev_name +
                                          activation_quant_type + '_' +
                                          weight_quant_type, loss_v))
W
WangZhen 已提交
345

346 347 348 349 350 351 352 353 354 355 356
        test_data = next(test_reader())
        with fluid.program_guard(quantized_test_program):
            w_var = fluid.framework._get_var('conv2d_1.w_0.quantized',
                                             quantized_test_program)
        # Testing
        with fluid.scope_guard(scope):
            test_loss1, w_quant = exe.run(program=quantized_test_program,
                                          feed=feeder.feed(test_data),
                                          fetch_list=[loss, w_var])

        # Freeze graph for inference, but the weight of fc/conv is still float type.
357
        freeze_pass = QuantizationFreezePass(
358
            scope=scope, place=place, weight_quantize_type=weight_quant_type)
359
        freeze_pass.apply(test_graph)
Z
Zhen Wang 已提交
360
        if not for_ci:
Z
Zhen Wang 已提交
361 362 363 364
            marked_nodes = set()
            for op in test_graph.all_op_nodes():
                if op.name().find('quantize') > -1:
                    marked_nodes.add(op)
365 366
            test_graph.draw('.', 'test_freeze' + dev_name +
                            activation_quant_type + '_' + weight_quant_type,
Z
Zhen Wang 已提交
367
                            marked_nodes)
W
WangZhen 已提交
368

369 370 371 372 373 374
        server_program = test_graph.to_program()
        with fluid.scope_guard(scope):
            test_loss2, = exe.run(program=server_program,
                                  feed=feeder.feed(test_data),
                                  fetch_list=[loss])
        self.assertAlmostEqual(test_loss1, test_loss2, delta=5e-3)
Z
Zhen Wang 已提交
375
        if not for_ci:
376 377 378 379 380 381
            print(
                '{}: {}'.format('test_loss1' + dev_name + activation_quant_type
                                + '_' + weight_quant_type, test_loss1))
            print(
                '{}: {}'.format('test_loss2' + dev_name + activation_quant_type
                                + '_' + weight_quant_type, test_loss2))
382 383
        w_freeze = np.array(scope.find_var('conv2d_1.w_0').get_tensor())
        # Maybe failed, this is due to the calculation precision
384
        # self.assertAlmostEqual(np.sum(w_freeze), np.sum(w_quant))
Z
Zhen Wang 已提交
385
        if not for_ci:
386 387 388 389
            print('{}: {}'.format('w_freeze' + dev_name + activation_quant_type
                                  + '_' + weight_quant_type, np.sum(w_freeze)))
            print('{}: {}'.format('w_quant' + dev_name + activation_quant_type +
                                  '_' + weight_quant_type, np.sum(w_quant)))
390 391 392 393

        # Convert parameter to 8-bit.
        convert_int8_pass = ConvertToInt8Pass(scope=scope, place=place)
        convert_int8_pass.apply(test_graph)
Z
Zhen Wang 已提交
394
        if not for_ci:
Z
Zhen Wang 已提交
395 396 397 398
            marked_nodes = set()
            for op in test_graph.all_op_nodes():
                if op.name().find('quantize') > -1:
                    marked_nodes.add(op)
399 400
            test_graph.draw('.', 'test_int8' + dev_name + activation_quant_type
                            + '_' + weight_quant_type, marked_nodes)
401 402 403
        server_program_int8 = test_graph.to_program()
        # Save the 8-bit parameter and model file.
        with fluid.scope_guard(scope):
404 405 406 407
            fluid.io.save_inference_model(
                'server_int8' + dev_name + activation_quant_type + '_' +
                weight_quant_type, ['image', 'label'], [loss], exe,
                server_program_int8)
408 409
            # Test whether the 8-bit parameter and model file can be loaded successfully.
            [infer, feed, fetch] = fluid.io.load_inference_model(
410 411
                'server_int8' + dev_name + activation_quant_type + '_' +
                weight_quant_type, exe)
412 413 414 415
        # Check the loaded 8-bit weight.
        w_8bit = np.array(scope.find_var('conv2d_1.w_0.int8').get_tensor())
        self.assertEqual(w_8bit.dtype, np.int8)
        self.assertEqual(np.sum(w_8bit), np.sum(w_freeze))
Z
Zhen Wang 已提交
416
        if not for_ci:
417 418 419 420
            print('{}: {}'.format('w_8bit' + dev_name + activation_quant_type +
                                  '_' + weight_quant_type, np.sum(w_8bit)))
            print('{}: {}'.format('w_freeze' + dev_name + activation_quant_type
                                  + '_' + weight_quant_type, np.sum(w_freeze)))
421 422 423

        mobile_pass = TransformForMobilePass()
        mobile_pass.apply(test_graph)
Z
Zhen Wang 已提交
424
        if not for_ci:
Z
Zhen Wang 已提交
425 426 427 428
            marked_nodes = set()
            for op in test_graph.all_op_nodes():
                if op.name().find('quantize') > -1:
                    marked_nodes.add(op)
429 430
            test_graph.draw('.', 'test_mobile' + dev_name +
                            activation_quant_type + '_' + weight_quant_type,
Z
Zhen Wang 已提交
431
                            marked_nodes)
432 433 434

        mobile_program = test_graph.to_program()
        with fluid.scope_guard(scope):
435 436 437 438
            fluid.io.save_inference_model(
                'mobile_int8' + dev_name + activation_quant_type + '_' +
                weight_quant_type, ['image', 'label'], [loss], exe,
                mobile_program)
W
WangZhen 已提交
439

440
    def test_freeze_graph_cuda_dynamic(self):
W
WangZhen 已提交
441 442
        if fluid.core.is_compiled_with_cuda():
            with fluid.unique_name.guard():
Z
Zhen Wang 已提交
443
                self.freeze_graph(
444 445 446 447 448 449 450 451 452 453 454 455
                    True,
                    seed=1,
                    activation_quant_type='abs_max',
                    weight_quant_type='abs_max',
                    for_ci=True)
            with fluid.unique_name.guard():
                self.freeze_graph(
                    True,
                    seed=1,
                    activation_quant_type='abs_max',
                    weight_quant_type='channel_wise_abs_max',
                    for_ci=True)
W
WangZhen 已提交
456

457
    def test_freeze_graph_cpu_dynamic(self):
W
WangZhen 已提交
458
        with fluid.unique_name.guard():
459 460 461 462 463 464 465 466 467 468 469 470
            self.freeze_graph(
                False,
                seed=2,
                activation_quant_type='abs_max',
                weight_quant_type='abs_max',
                for_ci=True)
            self.freeze_graph(
                False,
                seed=2,
                activation_quant_type='abs_max',
                weight_quant_type='channel_wise_abs_max',
                for_ci=True)
W
WangZhen 已提交
471

472
    def test_freeze_graph_cuda_static(self):
W
WangZhen 已提交
473 474
        if fluid.core.is_compiled_with_cuda():
            with fluid.unique_name.guard():
Z
Zhen Wang 已提交
475
                self.freeze_graph(
476 477 478 479 480 481 482 483 484 485 486
                    True,
                    seed=1,
                    activation_quant_type='range_abs_max',
                    weight_quant_type='abs_max',
                    for_ci=True)
                self.freeze_graph(
                    True,
                    seed=1,
                    activation_quant_type='moving_average_abs_max',
                    weight_quant_type='abs_max',
                    for_ci=True)
487 488 489
                self.freeze_graph(
                    True,
                    seed=1,
490 491 492 493 494 495 496 497 498
                    activation_quant_type='range_abs_max',
                    weight_quant_type='channel_wise_abs_max',
                    for_ci=True)
                self.freeze_graph(
                    True,
                    seed=1,
                    activation_quant_type='moving_average_abs_max',
                    weight_quant_type='channel_wise_abs_max',
                    for_ci=True)
W
WangZhen 已提交
499

500
    def test_freeze_graph_cpu_static(self):
W
WangZhen 已提交
501
        with fluid.unique_name.guard():
Z
Zhen Wang 已提交
502
            self.freeze_graph(
503 504 505 506 507 508 509 510 511 512 513 514 515 516 517 518 519
                False,
                seed=2,
                activation_quant_type='range_abs_max',
                weight_quant_type='abs_max',
                for_ci=True)
            self.freeze_graph(
                False,
                seed=2,
                activation_quant_type='moving_average_abs_max',
                weight_quant_type='abs_max',
                for_ci=True)
            self.freeze_graph(
                False,
                seed=2,
                activation_quant_type='range_abs_max',
                weight_quant_type='channel_wise_abs_max',
                for_ci=True)
520
            self.freeze_graph(
521 522
                False,
                seed=2,
523 524 525
                activation_quant_type='moving_average_abs_max',
                weight_quant_type='channel_wise_abs_max',
                for_ci=True)
W
WangZhen 已提交
526 527


528 529 530 531 532 533 534 535 536 537 538 539 540 541 542 543 544 545
def quant_dequant_residual_block(num, quant_skip_pattern=None):
    def conv_bn_layer(input,
                      ch_out,
                      filter_size,
                      stride,
                      padding,
                      act='relu',
                      bias_attr=False):
        tmp = fluid.layers.conv2d(
            input=input,
            filter_size=filter_size,
            num_filters=ch_out,
            stride=stride,
            padding=padding,
            act=None,
            bias_attr=bias_attr)
        return fluid.layers.batch_norm(input=tmp, act=act)

546 547 548
    data1 = fluid.layers.data(name='image', shape=[1, 32, 32], dtype='float32')
    data2 = fluid.layers.data(
        name='matmul_input', shape=[16, 32, 32], dtype='float32')
549
    label = fluid.layers.data(name='label', shape=[1], dtype='int64')
550
    hidden = data1
551 552 553 554
    for _ in six.moves.xrange(num):
        conv = conv_bn_layer(hidden, 16, 3, 1, 1, act=None, bias_attr=True)
        short = conv_bn_layer(hidden, 16, 1, 1, 0, act=None)
        hidden = fluid.layers.elementwise_add(x=conv, y=short, act='relu')
555
    hidden = fluid.layers.matmul(hidden, data2, True, True)
556
    if isinstance(quant_skip_pattern, str):
557 558 559 560 561 562 563
        with fluid.name_scope(quant_skip_pattern):
            pool1 = fluid.layers.pool2d(
                input=hidden, pool_size=2, pool_type='avg', pool_stride=2)
            pool2 = fluid.layers.pool2d(
                input=hidden, pool_size=2, pool_type='max', pool_stride=2)
            pool_add = fluid.layers.elementwise_add(
                x=pool1, y=pool2, act='relu')
564 565 566 567 568 569 570 571 572 573 574 575
    elif isinstance(quant_skip_pattern, list):
        assert len(
            quant_skip_pattern
        ) > 1, 'test config error: the len of quant_skip_pattern list should be greater than 1.'
        with fluid.name_scope(quant_skip_pattern[0]):
            pool1 = fluid.layers.pool2d(
                input=hidden, pool_size=2, pool_type='avg', pool_stride=2)
            pool2 = fluid.layers.pool2d(
                input=hidden, pool_size=2, pool_type='max', pool_stride=2)
        with fluid.name_scope(quant_skip_pattern[1]):
            pool_add = fluid.layers.elementwise_add(
                x=pool1, y=pool2, act='relu')
576 577 578 579 580 581 582 583 584 585 586 587
    else:
        pool1 = fluid.layers.pool2d(
            input=hidden, pool_size=2, pool_type='avg', pool_stride=2)
        pool2 = fluid.layers.pool2d(
            input=hidden, pool_size=2, pool_type='max', pool_stride=2)
        pool_add = fluid.layers.elementwise_add(x=pool1, y=pool2, act='relu')
    fc = fluid.layers.fc(input=pool_add, size=10)
    loss = fluid.layers.cross_entropy(input=fc, label=label)
    loss = fluid.layers.mean(loss)
    return loss


588 589 590 591 592
class TestAddQuantDequantPass(unittest.TestCase):
    def setUp(self):
        self._target_ops = {'elementwise_add', 'pool2d'}
        self._target_grad_ops = {'elementwise_add_grad', 'pool2d_grad'}

593
    def check_graph(self, graph, skip_pattern=None):
594 595 596
        ops = graph.all_op_nodes()
        for op_node in ops:
            if op_node.name() in self._target_ops:
597 598 599 600 601 602 603 604 605
                user_skipped = False
                if isinstance(skip_pattern, list):
                    user_skipped = op_node.op().has_attr("op_namescope") and \
                                   any(pattern in op_node.op().attr("op_namescope") for pattern in skip_pattern)
                elif isinstance(skip_pattern, str):
                    user_skipped = op_node.op().has_attr("op_namescope") and \
                                   op_node.op().attr("op_namescope").find(skip_pattern) != -1

                if user_skipped:
606 607
                    continue

608 609 610 611 612 613 614 615 616 617 618 619 620
                in_nodes_all_not_persistable = True
                for input_name in op_node.input_arg_names():
                    in_node = graph._find_node_by_name(op_node.inputs,
                                                       input_name)
                    in_nodes_all_not_persistable = (
                        in_nodes_all_not_persistable and
                        not in_node.persistable())
                if not in_nodes_all_not_persistable:
                    continue
                input_names = op_node.input_arg_names()
                for input_name in input_names:
                    self.assertTrue(input_name.endswith('.quant_dequant'))

621 622 623 624
    def residual_block_quant(self,
                             quantizable_op_type,
                             skip_pattern=None,
                             for_ci=True):
625 626 627
        main = fluid.Program()
        startup = fluid.Program()
        with fluid.program_guard(main, startup):
628
            loss = quant_dequant_residual_block(2, skip_pattern)
629 630 631 632 633
            opt = fluid.optimizer.Adam(learning_rate=0.001)
            opt.minimize(loss)
        place = fluid.CPUPlace()
        graph = IrGraph(core.Graph(main.desc), for_test=False)
        add_quant_dequant_pass = AddQuantDequantPass(
634 635 636 637
            scope=fluid.global_scope(),
            place=place,
            skip_pattern=skip_pattern,
            quantizable_op_type=quantizable_op_type)
638 639 640 641 642 643 644
        add_quant_dequant_pass.apply(graph)
        if not for_ci:
            marked_nodes = set()
            for op in graph.all_op_nodes():
                if op.name().find('quant') > -1:
                    marked_nodes.add(op)
            graph.draw('.', 'add_quant_dequant_graph', marked_nodes)
645
        self.check_graph(graph, skip_pattern)
646 647 648 649 650 651 652 653 654 655
        program = graph.to_program()
        val_graph = IrGraph(core.Graph(program.desc), for_test=False)
        if not for_ci:
            val_marked_nodes = set()
            for op in val_graph.all_op_nodes():
                if op.name().find('quant') > -1:
                    val_marked_nodes.add(op)
            val_graph.draw('.', 'val_add_quant_dequant_graph', val_marked_nodes)

    def test_residual_block(self):
656 657 658
        quantizable_op_type = ['elementwise_add', 'pool2d', 'mul', 'matmul']
        self.residual_block_quant(
            quantizable_op_type, skip_pattern=None, for_ci=True)
659 660

    def test_residual_block_skip_pattern(self):
661 662 663
        quantizable_op_type = ['elementwise_add', 'pool2d', 'mul', 'matmul']
        self.residual_block_quant(
            quantizable_op_type, skip_pattern='skip_quant', for_ci=True)
664

665
    def test_residual_block_skip_pattern(self):
666
        quantizable_op_type = ['elementwise_add', 'pool2d', 'mul', 'matmul']
667
        self.residual_block_quant(
668 669 670
            quantizable_op_type,
            skip_pattern=['skip_quant1', 'skip_quant2'],
            for_ci=True)
671

672

W
WangZhen 已提交
673 674
if __name__ == '__main__':
    unittest.main()