|
| 1 | +# copyright (c) 2018 paddlepaddle authors. all rights reserved. |
| 2 | +# |
| 3 | +# licensed under the apache license, version 2.0 (the "license"); |
| 4 | +# you may not use this file except in compliance with the license. |
| 5 | +# you may obtain a copy of the license at |
| 6 | +# |
| 7 | +# http://www.apache.org/licenses/license-2.0 |
| 8 | +# |
| 9 | +# unless required by applicable law or agreed to in writing, software |
| 10 | +# distributed under the license is distributed on an "as is" basis, |
| 11 | +# without warranties or conditions of any kind, either express or implied. |
| 12 | +# see the license for the specific language governing permissions and |
| 13 | +# limitations under the license. |
| 14 | + |
| 15 | +from __future__ import print_function |
| 16 | + |
| 17 | +import os |
| 18 | +import numpy as np |
| 19 | +import random |
| 20 | +import shutil |
| 21 | +import time |
| 22 | +import unittest |
| 23 | +import logging |
| 24 | + |
| 25 | +import paddle |
| 26 | +import paddle.fluid as fluid |
| 27 | +from paddle.fluid.contrib.slim.quantization import ImperativeQuantAware |
| 28 | +from paddle.fluid.log_helper import get_logger |
| 29 | +from paddle.dataset.common import download |
| 30 | + |
| 31 | +from imperative_test_utils import fix_model_dict, ImperativeLenet |
| 32 | + |
| 33 | +os.environ["CPU_NUM"] = "1" |
| 34 | +if paddle.is_compiled_with_cuda(): |
| 35 | + fluid.set_flags({"FLAGS_cudnn_deterministic": True}) |
| 36 | + |
| 37 | +_logger = get_logger( |
| 38 | + __name__, logging.INFO, fmt='%(asctime)s-%(levelname)s: %(message)s') |
| 39 | + |
| 40 | + |
| 41 | +class TestImperativeQatAmp(unittest.TestCase): |
| 42 | + """ |
| 43 | + Test the combination of qat and amp. |
| 44 | + """ |
| 45 | + |
| 46 | + @classmethod |
| 47 | + def setUpClass(cls): |
| 48 | + timestamp = time.strftime('%Y-%m-%d-%H-%M-%S', time.localtime()) |
| 49 | + cls.root_path = os.path.join(os.getcwd(), |
| 50 | + "imperative_qat_amp_" + timestamp) |
| 51 | + cls.save_path = os.path.join(cls.root_path, "model") |
| 52 | + |
| 53 | + cls.download_path = 'dygraph_int8/download' |
| 54 | + cls.cache_folder = os.path.expanduser('~/.cache/paddle/dataset/' + |
| 55 | + cls.download_path) |
| 56 | + |
| 57 | + cls.lenet_url = "https://paddle-inference-dist.cdn.bcebos.com/int8/unittest_model_data/lenet_pretrained.tar.gz" |
| 58 | + cls.lenet_md5 = "953b802fb73b52fae42896e3c24f0afb" |
| 59 | + |
| 60 | + seed = 1 |
| 61 | + np.random.seed(seed) |
| 62 | + paddle.static.default_main_program().random_seed = seed |
| 63 | + paddle.static.default_startup_program().random_seed = seed |
| 64 | + |
| 65 | + @classmethod |
| 66 | + def tearDownClass(cls): |
| 67 | + try: |
| 68 | + shutil.rmtree(cls.root_path) |
| 69 | + except Exception as e: |
| 70 | + print("Failed to delete {} due to {}".format(cls.root_path, str(e))) |
| 71 | + |
| 72 | + def cache_unzipping(self, target_folder, zip_path): |
| 73 | + if not os.path.exists(target_folder): |
| 74 | + cmd = 'mkdir {0} && tar xf {1} -C {0}'.format(target_folder, |
| 75 | + zip_path) |
| 76 | + os.system(cmd) |
| 77 | + |
| 78 | + def download_model(self, data_url, data_md5, folder_name): |
| 79 | + download(data_url, self.download_path, data_md5) |
| 80 | + file_name = data_url.split('/')[-1] |
| 81 | + zip_path = os.path.join(self.cache_folder, file_name) |
| 82 | + print('Data is downloaded at {0}'.format(zip_path)) |
| 83 | + |
| 84 | + data_cache_folder = os.path.join(self.cache_folder, folder_name) |
| 85 | + self.cache_unzipping(data_cache_folder, zip_path) |
| 86 | + return data_cache_folder |
| 87 | + |
| 88 | + def set_vars(self): |
| 89 | + self.qat = ImperativeQuantAware() |
| 90 | + |
| 91 | + self.train_batch_num = 30 |
| 92 | + self.train_batch_size = 32 |
| 93 | + self.test_batch_num = 100 |
| 94 | + self.test_batch_size = 32 |
| 95 | + self.eval_acc_top1 = 0.99 |
| 96 | + |
| 97 | + def model_train(self, model, batch_num=-1, batch_size=32, use_amp=False): |
| 98 | + model.train() |
| 99 | + |
| 100 | + train_reader = paddle.batch( |
| 101 | + paddle.dataset.mnist.train(), batch_size=batch_size) |
| 102 | + adam = paddle.optimizer.Adam( |
| 103 | + learning_rate=0.001, parameters=model.parameters()) |
| 104 | + scaler = paddle.amp.GradScaler(init_loss_scaling=500) |
| 105 | + |
| 106 | + for batch_id, data in enumerate(train_reader()): |
| 107 | + x_data = np.array([x[0].reshape(1, 28, 28) |
| 108 | + for x in data]).astype('float32') |
| 109 | + y_data = np.array( |
| 110 | + [x[1] for x in data]).astype('int64').reshape(-1, 1) |
| 111 | + |
| 112 | + img = paddle.to_tensor(x_data) |
| 113 | + label = paddle.to_tensor(y_data) |
| 114 | + |
| 115 | + if use_amp: |
| 116 | + with paddle.amp.auto_cast(): |
| 117 | + out = model(img) |
| 118 | + acc = fluid.layers.accuracy(out, label) |
| 119 | + loss = fluid.layers.cross_entropy(out, label) |
| 120 | + avg_loss = fluid.layers.mean(loss) |
| 121 | + scaled_loss = scaler.scale(avg_loss) |
| 122 | + scaled_loss.backward() |
| 123 | + |
| 124 | + scaler.minimize(adam, scaled_loss) |
| 125 | + adam.clear_gradients() |
| 126 | + else: |
| 127 | + out = model(img) |
| 128 | + acc = fluid.layers.accuracy(out, label) |
| 129 | + loss = fluid.layers.cross_entropy(out, label) |
| 130 | + avg_loss = fluid.layers.mean(loss) |
| 131 | + avg_loss.backward() |
| 132 | + |
| 133 | + adam.minimize(avg_loss) |
| 134 | + model.clear_gradients() |
| 135 | + |
| 136 | + if batch_id % 100 == 0: |
| 137 | + _logger.info("Train | step {}: loss = {:}, acc= {:}".format( |
| 138 | + batch_id, avg_loss.numpy(), acc.numpy())) |
| 139 | + |
| 140 | + if batch_num > 0 and batch_id + 1 >= batch_num: |
| 141 | + break |
| 142 | + |
| 143 | + def model_test(self, model, batch_num=-1, batch_size=32, use_amp=False): |
| 144 | + model.eval() |
| 145 | + |
| 146 | + test_reader = paddle.batch( |
| 147 | + paddle.dataset.mnist.test(), batch_size=batch_size) |
| 148 | + |
| 149 | + acc_top1_list = [] |
| 150 | + for batch_id, data in enumerate(test_reader()): |
| 151 | + x_data = np.array([x[0].reshape(1, 28, 28) |
| 152 | + for x in data]).astype('float32') |
| 153 | + y_data = np.array( |
| 154 | + [x[1] for x in data]).astype('int64').reshape(-1, 1) |
| 155 | + |
| 156 | + img = paddle.to_tensor(x_data) |
| 157 | + label = paddle.to_tensor(y_data) |
| 158 | + |
| 159 | + with paddle.amp.auto_cast(use_amp): |
| 160 | + out = model(img) |
| 161 | + acc_top1 = fluid.layers.accuracy(input=out, label=label, k=1) |
| 162 | + acc_top5 = fluid.layers.accuracy(input=out, label=label, k=5) |
| 163 | + |
| 164 | + acc_top1_list.append(float(acc_top1.numpy())) |
| 165 | + if batch_id % 100 == 0: |
| 166 | + _logger.info("Test | At step {}: acc1 = {:}, acc5 = {:}".format( |
| 167 | + batch_id, acc_top1.numpy(), acc_top5.numpy())) |
| 168 | + |
| 169 | + if batch_num > 0 and batch_id + 1 >= batch_num: |
| 170 | + break |
| 171 | + |
| 172 | + acc_top1 = sum(acc_top1_list) / len(acc_top1_list) |
| 173 | + return acc_top1 |
| 174 | + |
| 175 | + def test_ptq(self): |
| 176 | + start_time = time.time() |
| 177 | + |
| 178 | + self.set_vars() |
| 179 | + |
| 180 | + params_path = self.download_model(self.lenet_url, self.lenet_md5, |
| 181 | + "lenet") |
| 182 | + params_path += "/lenet_pretrained/lenet.pdparams" |
| 183 | + |
| 184 | + with fluid.dygraph.guard(): |
| 185 | + model = ImperativeLenet() |
| 186 | + model_state_dict = paddle.load(params_path) |
| 187 | + model.set_state_dict(model_state_dict) |
| 188 | + |
| 189 | + _logger.info("Test fp32 model") |
| 190 | + fp32_acc_top1 = self.model_test(model, self.test_batch_num, |
| 191 | + self.test_batch_size) |
| 192 | + |
| 193 | + self.qat.quantize(model) |
| 194 | + |
| 195 | + use_amp = True |
| 196 | + self.model_train(model, self.train_batch_num, self.train_batch_size, |
| 197 | + use_amp) |
| 198 | + |
| 199 | + _logger.info("Test int8 model") |
| 200 | + int8_acc_top1 = self.model_test(model, self.test_batch_num, |
| 201 | + self.test_batch_size, use_amp) |
| 202 | + |
| 203 | + _logger.info('fp32_acc_top1: %f, int8_acc_top1: %f' % |
| 204 | + (fp32_acc_top1, int8_acc_top1)) |
| 205 | + self.assertTrue( |
| 206 | + int8_acc_top1 > fp32_acc_top1 - 0.01, |
| 207 | + msg='fp32_acc_top1: %f, int8_acc_top1: %f' % |
| 208 | + (fp32_acc_top1, int8_acc_top1)) |
| 209 | + |
| 210 | + input_spec = [ |
| 211 | + paddle.static.InputSpec( |
| 212 | + shape=[None, 1, 28, 28], dtype='float32') |
| 213 | + ] |
| 214 | + paddle.jit.save(layer=model, path=self.save_path, input_spec=input_spec) |
| 215 | + print('Quantized model saved in {%s}' % self.save_path) |
| 216 | + |
| 217 | + end_time = time.time() |
| 218 | + print("total time: %ss" % (end_time - start_time)) |
| 219 | + |
| 220 | + |
| 221 | +if __name__ == '__main__': |
| 222 | + unittest.main() |
0 commit comments