program.py 28.5 KB
Newer Older
L
LDOUBLEV 已提交
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24
# 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 absolute_import
from __future__ import division
from __future__ import print_function

from argparse import ArgumentParser, RawDescriptionHelpFormatter
import sys
import yaml
import os
from ppocr.utils.utility import create_module
from ppocr.utils.utility import initial_logger
L
licx 已提交
25

L
LDOUBLEV 已提交
26 27 28 29 30 31 32
logger = initial_logger()

import paddle.fluid as fluid
import time
from ppocr.utils.stats import TrainingStats
from eval_utils.eval_det_utils import eval_det_run
from eval_utils.eval_rec_utils import eval_rec_run
W
WenmuZhou 已提交
33
from eval_utils.eval_cls_utils import eval_cls_run
L
LDOUBLEV 已提交
34 35
from ppocr.utils.save_load import save_model
import numpy as np
T
tink2123 已提交
36
from ppocr.utils.character import cal_predicts_accuracy, cal_predicts_accuracy_srn, CharacterOps
L
LDOUBLEV 已提交
37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79


class ArgsParser(ArgumentParser):
    def __init__(self):
        super(ArgsParser, self).__init__(
            formatter_class=RawDescriptionHelpFormatter)
        self.add_argument("-c", "--config", help="configuration file to use")
        self.add_argument(
            "-o", "--opt", nargs='+', help="set configuration options")

    def parse_args(self, argv=None):
        args = super(ArgsParser, self).parse_args(argv)
        assert args.config is not None, \
            "Please specify --config=configure_file_path."
        args.opt = self._parse_opt(args.opt)
        return args

    def _parse_opt(self, opts):
        config = {}
        if not opts:
            return config
        for s in opts:
            s = s.strip()
            k, v = s.split('=')
            config[k] = yaml.load(v, Loader=yaml.Loader)
        return config


class AttrDict(dict):
    """Single level attribute dict, NOT recursive"""

    def __init__(self, **kwargs):
        super(AttrDict, self).__init__()
        super(AttrDict, self).update(kwargs)

    def __getattr__(self, key):
        if key in self:
            return self[key]
        raise AttributeError("object has no attribute '{}'".format(key))


global_config = AttrDict()

农夫三拳_'s avatar
农夫三拳_ 已提交
80 81
default_config = {'Global': {'debug': False, }}

L
LDOUBLEV 已提交
82 83 84 85 86 87 88 89

def load_config(file_path):
    """
    Load config from yml/yaml file.
    Args:
        file_path (str): Path of the config file to be loaded.
    Returns: global config
    """
农夫三拳_'s avatar
农夫三拳_ 已提交
90
    merge_config(default_config)
L
LDOUBLEV 已提交
91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117
    _, ext = os.path.splitext(file_path)
    assert ext in ['.yml', '.yaml'], "only support yaml files for now"
    merge_config(yaml.load(open(file_path), Loader=yaml.Loader))
    assert "reader_yml" in global_config['Global'],\
        "absence reader_yml in global"
    reader_file_path = global_config['Global']['reader_yml']
    _, ext = os.path.splitext(reader_file_path)
    assert ext in ['.yml', '.yaml'], "only support yaml files for reader"
    merge_config(yaml.load(open(reader_file_path), Loader=yaml.Loader))
    return global_config


def merge_config(config):
    """
    Merge config into global config.
    Args:
        config (dict): Config to be merged.
    Returns: global config
    """
    for key, value in config.items():
        if "." not in key:
            if isinstance(value, dict) and key in global_config:
                global_config[key].update(value)
            else:
                global_config[key] = value
        else:
            sub_keys = key.split('.')
T
tink2123 已提交
118 119 120 121
            assert (
                sub_keys[0] in global_config
            ), "the sub_keys can only be one of global_config: {}, but get: {}, please check your running command".format(
                global_config.keys(), sub_keys[0])
L
LDOUBLEV 已提交
122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152
            cur = global_config[sub_keys[0]]
            for idx, sub_key in enumerate(sub_keys[1:]):
                assert (sub_key in cur)
                if idx == len(sub_keys) - 2:
                    cur[sub_key] = value
                else:
                    cur = cur[sub_key]


def check_gpu(use_gpu):
    """
    Log error and exit when set use_gpu=true in paddlepaddle
    cpu version.
    """
    err = "Config use_gpu cannot be set as true while you are " \
          "using paddlepaddle cpu version ! \nPlease try: \n" \
          "\t1. Install paddlepaddle-gpu to run model on GPU \n" \
          "\t2. Set use_gpu as false in config file to run " \
          "model on CPU"

    try:
        if use_gpu and not fluid.is_compiled_with_cuda():
            logger.error(err)
            sys.exit(1)
    except Exception as e:
        pass


def build(config, main_prog, startup_prog, mode):
    """
    Build a program using a model and an optimizer
T
tink2123 已提交
153 154 155 156
        1. create a dataloader
        2. create a model
        3. create fetchs
        4. create an optimizer
L
LDOUBLEV 已提交
157 158 159 160
    Args:
        config(dict): config
        main_prog(): main program
        startup_prog(): startup program
T
tink2123 已提交
161
        mode(str): train or valid
L
LDOUBLEV 已提交
162 163
    Returns:
        dataloader(): a bridge between the model and the data
T
tink2123 已提交
164 165 166
        fetch_name_list(dict): dict of model outputs(included loss and measures)
        fetch_varname_list(list): list of outputs' varname
        opt_loss_name(str): name of loss
L
LDOUBLEV 已提交
167 168 169 170 171 172 173 174 175
    """
    with fluid.program_guard(main_prog, startup_prog):
        with fluid.unique_name.guard():
            func_infor = config['Architecture']['function']
            model = create_module(func_infor)(params=config)
            dataloader, outputs = model(mode=mode)
            fetch_name_list = list(outputs.keys())
            fetch_varname_list = [outputs[v].name for v in fetch_name_list]
            opt_loss_name = None
T
tink2123 已提交
176 177 178
            model_average = None
            img_loss_name = None
            word_loss_name = None
L
LDOUBLEV 已提交
179 180
            if mode == "train":
                opt_loss = outputs['total_loss']
T
tink2123 已提交
181 182 183 184 185
                # srn loss
                #img_loss = outputs['img_loss']
                #word_loss = outputs['word_loss']
                #img_loss_name = img_loss.name
                #word_loss_name = word_loss.name
L
LDOUBLEV 已提交
186 187 188 189 190 191 192
                opt_params = config['Optimizer']
                optimizer = create_module(opt_params['function'])(opt_params)
                optimizer.minimize(opt_loss)
                opt_loss_name = opt_loss.name
                global_lr = optimizer._global_learning_rate()
                fetch_name_list.insert(0, "lr")
                fetch_varname_list.insert(0, global_lr.name)
T
tink2123 已提交
193 194 195 196 197 198 199 200
                if "loss_type" in config["Global"]:
                    if config['Global']["loss_type"] == 'srn':
                        model_average = fluid.optimizer.ModelAverage(
                            config['Global']['average_window'],
                            min_average_window=config['Global'][
                                'min_average_window'],
                            max_average_window=config['Global'][
                                'max_average_window'])
T
tink2123 已提交
201

T
tink2123 已提交
202 203
    return (dataloader, fetch_name_list, fetch_varname_list, opt_loss_name,
            model_average)
L
LDOUBLEV 已提交
204 205 206 207


def build_export(config, main_prog, startup_prog):
    """
208 209 210 211 212 213 214 215 216
    Build input and output for exporting a checkpoints model to an inference model
    Args:
        config(dict): config
        main_prog(): main program
        startup_prog(): startup program
    Returns:
        feeded_var_names(list[str]): var names of input for exported inference model
        target_vars(list[Variable]): output vars for exported inference model
        fetches_var_name: dict of checkpoints model outputs(included loss and measures)
L
LDOUBLEV 已提交
217 218 219 220 221
    """
    with fluid.program_guard(main_prog, startup_prog):
        with fluid.unique_name.guard():
            func_infor = config['Architecture']['function']
            model = create_module(func_infor)(params=config)
T
tink2123 已提交
222 223
            algorithm = config['Global']['algorithm']
            if algorithm == "SRN":
T
tink2123 已提交
224 225 226
                image, others, outputs = model(mode='export')
            else:
                image, outputs = model(mode='export')
227
            fetches_var_name = sorted([name for name in outputs.keys()])
D
dyning 已提交
228
            fetches_var = [outputs[name] for name in fetches_var_name]
T
tink2123 已提交
229
    if algorithm == "SRN":
T
tink2123 已提交
230 231 232 233 234
        others_var_names = sorted([name for name in others.keys()])
        feeded_var_names = [image.name] + others_var_names
    else:
        feeded_var_names = [image.name]

L
LDOUBLEV 已提交
235 236 237 238
    target_vars = fetches_var
    return feeded_var_names, target_vars, fetches_var_name


B
baiyfbupt 已提交
239
def create_multi_devices_program(program, loss_var_name, for_quant=False):
L
LDOUBLEV 已提交
240 241 242
    build_strategy = fluid.BuildStrategy()
    build_strategy.memory_optimize = False
    build_strategy.enable_inplace = True
B
baiyfbupt 已提交
243 244
    if for_quant:
        build_strategy.fuse_all_reduce_ops = False
B
baiyfbupt 已提交
245 246
    else:
        program = fluid.CompiledProgram(program)
L
LDOUBLEV 已提交
247 248
    exec_strategy = fluid.ExecutionStrategy()
    exec_strategy.num_iteration_per_drop_scope = 1
B
baiyfbupt 已提交
249
    compile_program = program.with_data_parallel(
L
LDOUBLEV 已提交
250 251 252 253 254 255
        loss_name=loss_var_name,
        build_strategy=build_strategy,
        exec_strategy=exec_strategy)
    return compile_program


Y
yukavio 已提交
256 257 258 259
def train_eval_det_run(config,
                       exe,
                       train_info_dict,
                       eval_info_dict,
B
baiyfbupt 已提交
260
                       is_slim=None):
T
tink2123 已提交
261 262 263 264 265 266 267 268
    """
    Feed data to the model and fetch the measures and loss for detection
    Args:
        config: config
        exe:
        train_info_dict: information dict for training
        eval_info_dict: information dict for evaluation
    """
L
LDOUBLEV 已提交
269 270 271 272 273
    train_batch_id = 0
    log_smooth_window = config['Global']['log_smooth_window']
    epoch_num = config['Global']['epoch_num']
    print_batch_step = config['Global']['print_batch_step']
    eval_batch_step = config['Global']['eval_batch_step']
L
LDOUBLEV 已提交
274 275 276 277 278 279 280
    start_eval_step = 0
    if type(eval_batch_step) == list and len(eval_batch_step) >= 2:
        start_eval_step = eval_batch_step[0]
        eval_batch_step = eval_batch_step[1]
        logger.info(
            "During the training process, after the {}th iteration, an evaluation is run every {} iterations".
            format(start_eval_step, eval_batch_step))
L
LDOUBLEV 已提交
281 282
    save_epoch_step = config['Global']['save_epoch_step']
    save_model_dir = config['Global']['save_model_dir']
283 284
    if not os.path.exists(save_model_dir):
        os.makedirs(save_model_dir)
L
LDOUBLEV 已提交
285 286 287 288 289 290 291 292 293 294 295 296 297 298 299 300 301 302 303 304 305 306 307
    train_stats = TrainingStats(log_smooth_window,
                                train_info_dict['fetch_name_list'])
    best_eval_hmean = -1
    best_batch_id = 0
    best_epoch = 0
    train_loader = train_info_dict['reader']
    for epoch in range(epoch_num):
        train_loader.start()
        try:
            while True:
                t1 = time.time()
                train_outs = exe.run(
                    program=train_info_dict['compile_program'],
                    fetch_list=train_info_dict['fetch_varname_list'],
                    return_numpy=False)
                stats = {}
                for tno in range(len(train_outs)):
                    fetch_name = train_info_dict['fetch_name_list'][tno]
                    fetch_value = np.mean(np.array(train_outs[tno]))
                    stats[fetch_name] = fetch_value
                t2 = time.time()
                train_batch_elapse = t2 - t1
                train_stats.update(stats)
L
LDOUBLEV 已提交
308
                if train_batch_id > 0 and train_batch_id  \
L
LDOUBLEV 已提交
309 310 311 312 313 314
                    % print_batch_step == 0:
                    logs = train_stats.log()
                    strs = 'epoch: {}, iter: {}, {}, time: {:.3f}'.format(
                        epoch, train_batch_id, logs, train_batch_elapse)
                    logger.info(strs)

L
LDOUBLEV 已提交
315 316
                if train_batch_id > start_eval_step and\
                    (train_batch_id - start_eval_step) % eval_batch_step == 0:
L
LDOUBLEV 已提交
317 318 319 320 321 322 323
                    metrics = eval_det_run(exe, config, eval_info_dict, "eval")
                    hmean = metrics['hmean']
                    if hmean >= best_eval_hmean:
                        best_eval_hmean = hmean
                        best_batch_id = train_batch_id
                        best_epoch = epoch
                        save_path = save_model_dir + "/best_accuracy"
B
baiyfbupt 已提交
324
                        if is_slim is None:
Y
yukavio 已提交
325 326
                            save_model(train_info_dict['train_program'],
                                       save_path)
B
baiyfbupt 已提交
327 328 329 330 331 332 333 334
                        else:
                            import paddleslim as slim
                            if is_slim == "prune":
                                slim.prune.save_model(
                                    exe, train_info_dict['train_program'],
                                    save_path)
                            elif is_slim == "quant":
                                save_model(eval_info_dict['program'], save_path)
B
baiyfbupt 已提交
335 336 337 338
                            else:
                                raise ValueError(
                                    "Only quant and prune are supported currently. But received {}".
                                    format(is_slim))
L
LDOUBLEV 已提交
339 340 341 342 343 344 345 346
                    strs = 'Test iter: {}, metrics:{}, best_hmean:{:.6f}, best_epoch:{}, best_batch_id:{}'.format(
                        train_batch_id, metrics, best_eval_hmean, best_epoch,
                        best_batch_id)
                    logger.info(strs)
                train_batch_id += 1

        except fluid.core.EOFException:
            train_loader.reset()
T
tink2123 已提交
347
        if epoch == 0 and save_epoch_step == 1:
T
tink2123 已提交
348
            save_path = save_model_dir + "/iter_epoch_0"
B
baiyfbupt 已提交
349
            if is_slim is None:
Y
yukavio 已提交
350
                save_model(train_info_dict['train_program'], save_path)
B
baiyfbupt 已提交
351 352 353 354 355 356 357
            else:
                import paddleslim as slim
                if is_slim == "prune":
                    slim.prune.save_model(exe, train_info_dict['train_program'],
                                          save_path)
                elif is_slim == "quant":
                    save_model(eval_info_dict['program'], save_path)
B
baiyfbupt 已提交
358 359 360 361
                else:
                    raise ValueError(
                        "Only quant and prune are supported currently. But received {}".
                        format(is_slim))
L
LDOUBLEV 已提交
362 363
        if epoch > 0 and epoch % save_epoch_step == 0:
            save_path = save_model_dir + "/iter_epoch_%d" % (epoch)
B
baiyfbupt 已提交
364
            if is_slim is None:
Y
yukavio 已提交
365
                save_model(train_info_dict['train_program'], save_path)
B
baiyfbupt 已提交
366 367 368 369 370 371 372
            else:
                import paddleslim as slim
                if is_slim == "prune":
                    slim.prune.save_model(exe, train_info_dict['train_program'],
                                          save_path)
                elif is_slim == "quant":
                    save_model(eval_info_dict['program'], save_path)
B
baiyfbupt 已提交
373 374 375 376
                else:
                    raise ValueError(
                        "Only quant and prune are supported currently. But received {}".
                        format(is_slim))
L
LDOUBLEV 已提交
377 378 379
    return


B
baiyfbupt 已提交
380 381 382 383 384
def train_eval_rec_run(config,
                       exe,
                       train_info_dict,
                       eval_info_dict,
                       is_slim=None):
T
tink2123 已提交
385 386 387 388 389 390 391 392
    """
    Feed data to the model and fetch the measures and loss for recognition
    Args:
        config: config
        exe:
        train_info_dict: information dict for training
        eval_info_dict: information dict for evaluation
    """
L
LDOUBLEV 已提交
393 394 395 396 397
    train_batch_id = 0
    log_smooth_window = config['Global']['log_smooth_window']
    epoch_num = config['Global']['epoch_num']
    print_batch_step = config['Global']['print_batch_step']
    eval_batch_step = config['Global']['eval_batch_step']
L
LDOUBLEV 已提交
398 399 400 401 402 403 404
    start_eval_step = 0
    if type(eval_batch_step) == list and len(eval_batch_step) >= 2:
        start_eval_step = eval_batch_step[0]
        eval_batch_step = eval_batch_step[1]
        logger.info(
            "During the training process, after the {}th iteration, an evaluation is run every {} iterations".
            format(start_eval_step, eval_batch_step))
L
LDOUBLEV 已提交
405 406
    save_epoch_step = config['Global']['save_epoch_step']
    save_model_dir = config['Global']['save_model_dir']
L
LDOUBLEV 已提交
407
    if not os.path.exists(save_model_dir):
L
LDOUBLEV 已提交
408
        os.makedirs(save_model_dir)
L
LDOUBLEV 已提交
409 410 411 412 413 414 415 416 417 418 419 420 421 422 423 424 425 426 427 428 429 430 431 432 433
    train_stats = TrainingStats(log_smooth_window, ['loss', 'acc'])
    best_eval_acc = -1
    best_batch_id = 0
    best_epoch = 0
    train_loader = train_info_dict['reader']
    for epoch in range(epoch_num):
        train_loader.start()
        try:
            while True:
                t1 = time.time()
                train_outs = exe.run(
                    program=train_info_dict['compile_program'],
                    fetch_list=train_info_dict['fetch_varname_list'],
                    return_numpy=False)
                fetch_map = dict(
                    zip(train_info_dict['fetch_name_list'],
                        range(len(train_outs))))

                loss = np.mean(np.array(train_outs[fetch_map['total_loss']]))
                lr = np.mean(np.array(train_outs[fetch_map['lr']]))
                preds_idx = fetch_map['decoded_out']
                preds = np.array(train_outs[preds_idx])
                labels_idx = fetch_map['label']
                labels = np.array(train_outs[labels_idx])

T
tink2123 已提交
434 435 436 437 438 439 440 441 442 443 444
                if config['Global']['loss_type'] != 'srn':
                    preds_lod = train_outs[preds_idx].lod()[0]
                    labels_lod = train_outs[labels_idx].lod()[0]

                    acc, acc_num, img_num = cal_predicts_accuracy(
                        config['Global']['char_ops'], preds, preds_lod, labels,
                        labels_lod)
                else:
                    acc, acc_num, img_num = cal_predicts_accuracy_srn(
                        config['Global']['char_ops'], preds, labels,
                        config['Global']['max_text_length'])
L
LDOUBLEV 已提交
445 446 447 448
                t2 = time.time()
                train_batch_elapse = t2 - t1
                stats = {'loss': loss, 'acc': acc}
                train_stats.update(stats)
L
update  
LDOUBLEV 已提交
449
                if train_batch_id > start_eval_step and (train_batch_id - start_eval_step) \
L
LDOUBLEV 已提交
450 451 452 453 454 455 456 457
                    % print_batch_step == 0:
                    logs = train_stats.log()
                    strs = 'epoch: {}, iter: {}, lr: {:.6f}, {}, time: {:.3f}'.format(
                        epoch, train_batch_id, lr, logs, train_batch_elapse)
                    logger.info(strs)

                if train_batch_id > 0 and\
                    train_batch_id % eval_batch_step == 0:
T
tink2123 已提交
458 459 460
                    model_average = train_info_dict['model_average']
                    if model_average != None:
                        model_average.apply(exe)
L
LDOUBLEV 已提交
461 462 463 464 465 466 467 468
                    metrics = eval_rec_run(exe, config, eval_info_dict, "eval")
                    eval_acc = metrics['avg_acc']
                    eval_sample_num = metrics['total_sample_num']
                    if eval_acc > best_eval_acc:
                        best_eval_acc = eval_acc
                        best_batch_id = train_batch_id
                        best_epoch = epoch
                        save_path = save_model_dir + "/best_accuracy"
B
baiyfbupt 已提交
469 470 471 472 473 474 475 476 477 478 479
                        if is_slim is None:
                            save_model(train_info_dict['train_program'],
                                       save_path)
                        else:
                            import paddleslim as slim
                            if is_slim == "prune":
                                slim.prune.save_model(
                                    exe, train_info_dict['train_program'],
                                    save_path)
                            elif is_slim == "quant":
                                save_model(eval_info_dict['program'], save_path)
B
baiyfbupt 已提交
480 481 482 483
                            else:
                                raise ValueError(
                                    "Only quant and prune are supported currently. But received {}".
                                    format(is_slim))
L
LDOUBLEV 已提交
484 485 486 487 488 489 490 491
                    strs = 'Test iter: {}, acc:{:.6f}, best_acc:{:.6f}, best_epoch:{}, best_batch_id:{}, eval_sample_num:{}'.format(
                        train_batch_id, eval_acc, best_eval_acc, best_epoch,
                        best_batch_id, eval_sample_num)
                    logger.info(strs)
                train_batch_id += 1

        except fluid.core.EOFException:
            train_loader.reset()
T
tink2123 已提交
492
        if epoch == 0 and save_epoch_step == 1:
T
tink2123 已提交
493
            save_path = save_model_dir + "/iter_epoch_0"
B
baiyfbupt 已提交
494 495 496 497 498 499 500 501 502
            if is_slim is None:
                save_model(train_info_dict['train_program'], save_path)
            else:
                import paddleslim as slim
                if is_slim == "prune":
                    slim.prune.save_model(exe, train_info_dict['train_program'],
                                          save_path)
                elif is_slim == "quant":
                    save_model(eval_info_dict['program'], save_path)
B
baiyfbupt 已提交
503 504 505 506
                else:
                    raise ValueError(
                        "Only quant and prune are supported currently. But received {}".
                        format(is_slim))
L
LDOUBLEV 已提交
507 508
        if epoch > 0 and epoch % save_epoch_step == 0:
            save_path = save_model_dir + "/iter_epoch_%d" % (epoch)
B
baiyfbupt 已提交
509 510 511 512 513 514 515 516 517
            if is_slim is None:
                save_model(train_info_dict['train_program'], save_path)
            else:
                import paddleslim as slim
                if is_slim == "prune":
                    slim.prune.save_model(exe, train_info_dict['train_program'],
                                          save_path)
                elif is_slim == "quant":
                    save_model(eval_info_dict['program'], save_path)
B
baiyfbupt 已提交
518 519 520 521
                else:
                    raise ValueError(
                        "Only quant and prune are supported currently. But received {}".
                        format(is_slim))
L
LDOUBLEV 已提交
522
    return
L
licx 已提交
523

T
tink2123 已提交
524

B
baiyfbupt 已提交
525 526 527 528 529
def train_eval_cls_run(config,
                       exe,
                       train_info_dict,
                       eval_info_dict,
                       is_slim=None):
W
WenmuZhou 已提交
530 531 532 533 534 535 536 537 538 539 540 541 542 543 544 545 546 547 548 549 550 551 552 553 554 555 556 557 558 559 560 561 562 563 564 565 566 567 568 569 570 571 572 573 574 575 576 577 578 579 580 581 582 583 584 585 586 587 588 589 590 591
    train_batch_id = 0
    log_smooth_window = config['Global']['log_smooth_window']
    epoch_num = config['Global']['epoch_num']
    print_batch_step = config['Global']['print_batch_step']
    eval_batch_step = config['Global']['eval_batch_step']
    start_eval_step = 0
    if type(eval_batch_step) == list and len(eval_batch_step) >= 2:
        start_eval_step = eval_batch_step[0]
        eval_batch_step = eval_batch_step[1]
        logger.info(
            "During the training process, after the {}th iteration, an evaluation is run every {} iterations".
            format(start_eval_step, eval_batch_step))
    save_epoch_step = config['Global']['save_epoch_step']
    save_model_dir = config['Global']['save_model_dir']
    if not os.path.exists(save_model_dir):
        os.makedirs(save_model_dir)
    train_stats = TrainingStats(log_smooth_window, ['loss', 'acc'])
    best_eval_acc = -1
    best_batch_id = 0
    best_epoch = 0
    train_loader = train_info_dict['reader']
    for epoch in range(epoch_num):
        train_loader.start()
        try:
            while True:
                t1 = time.time()
                train_outs = exe.run(
                    program=train_info_dict['compile_program'],
                    fetch_list=train_info_dict['fetch_varname_list'],
                    return_numpy=False)
                fetch_map = dict(
                    zip(train_info_dict['fetch_name_list'],
                        range(len(train_outs))))

                loss = np.mean(np.array(train_outs[fetch_map['total_loss']]))
                lr = np.mean(np.array(train_outs[fetch_map['lr']]))
                acc = np.mean(np.array(train_outs[fetch_map['acc']]))

                t2 = time.time()
                train_batch_elapse = t2 - t1
                stats = {'loss': loss, 'acc': acc}
                train_stats.update(stats)
                if train_batch_id > start_eval_step and (train_batch_id - start_eval_step) \
                    % print_batch_step == 0:
                    logs = train_stats.log()
                    strs = 'epoch: {}, iter: {}, lr: {:.6f}, {}, time: {:.3f}'.format(
                        epoch, train_batch_id, lr, logs, train_batch_elapse)
                    logger.info(strs)

                if train_batch_id > 0 and\
                    train_batch_id % eval_batch_step == 0:
                    model_average = train_info_dict['model_average']
                    if model_average != None:
                        model_average.apply(exe)
                    metrics = eval_cls_run(exe, eval_info_dict)
                    eval_acc = metrics['avg_acc']
                    eval_sample_num = metrics['total_sample_num']
                    if eval_acc > best_eval_acc:
                        best_eval_acc = eval_acc
                        best_batch_id = train_batch_id
                        best_epoch = epoch
                        save_path = save_model_dir + "/best_accuracy"
B
baiyfbupt 已提交
592 593 594 595 596 597 598 599 600 601 602
                        if is_slim is None:
                            save_model(train_info_dict['train_program'],
                                       save_path)
                        else:
                            import paddleslim as slim
                            if is_slim == "prune":
                                slim.prune.save_model(
                                    exe, train_info_dict['train_program'],
                                    save_path)
                            elif is_slim == "quant":
                                save_model(eval_info_dict['program'], save_path)
B
baiyfbupt 已提交
603 604 605 606
                            else:
                                raise ValueError(
                                    "Only quant and prune are supported currently. But received {}".
                                    format(is_slim))
W
WenmuZhou 已提交
607 608 609 610 611 612 613 614 615 616
                    strs = 'Test iter: {}, acc:{:.6f}, best_acc:{:.6f}, best_epoch:{}, best_batch_id:{}, eval_sample_num:{}'.format(
                        train_batch_id, eval_acc, best_eval_acc, best_epoch,
                        best_batch_id, eval_sample_num)
                    logger.info(strs)
                train_batch_id += 1

        except fluid.core.EOFException:
            train_loader.reset()
        if epoch == 0 and save_epoch_step == 1:
            save_path = save_model_dir + "/iter_epoch_0"
B
baiyfbupt 已提交
617 618 619 620 621 622 623 624 625
            if is_slim is None:
                save_model(train_info_dict['train_program'], save_path)
            else:
                import paddleslim as slim
                if is_slim == "prune":
                    slim.prune.save_model(exe, train_info_dict['train_program'],
                                          save_path)
                elif is_slim == "quant":
                    save_model(eval_info_dict['program'], save_path)
B
baiyfbupt 已提交
626 627 628 629
                else:
                    raise ValueError(
                        "Only quant and prune are supported currently. But received {}".
                        format(is_slim))
W
WenmuZhou 已提交
630 631
        if epoch > 0 and epoch % save_epoch_step == 0:
            save_path = save_model_dir + "/iter_epoch_%d" % (epoch)
B
baiyfbupt 已提交
632 633 634 635 636 637 638 639 640
            if is_slim is None:
                save_model(train_info_dict['train_program'], save_path)
            else:
                import paddleslim as slim
                if is_slim == "prune":
                    slim.prune.save_model(exe, train_info_dict['train_program'],
                                          save_path)
                elif is_slim == "quant":
                    save_model(eval_info_dict['program'], save_path)
B
baiyfbupt 已提交
641 642 643 644
                else:
                    raise ValueError(
                        "Only quant and prune are supported currently. But received {}".
                        format(is_slim))
W
WenmuZhou 已提交
645 646 647
    return


L
licx 已提交
648
def preprocess():
649
    # load config from yml file
L
licx 已提交
650 651 652 653 654 655 656 657 658
    FLAGS = ArgsParser().parse_args()
    config = load_config(FLAGS.config)
    merge_config(FLAGS.opt)
    logger.info(config)

    # check if set use_gpu=True in paddlepaddle cpu version
    use_gpu = config['Global']['use_gpu']
    check_gpu(use_gpu)

659
    # check whether the set algorithm belongs to the supported algorithm list
L
licx 已提交
660
    alg = config['Global']['algorithm']
T
tink2123 已提交
661
    assert alg in [
W
WenmuZhou 已提交
662
        'EAST', 'DB', 'SAST', 'Rosetta', 'CRNN', 'STARNet', 'RARE', 'SRN', 'CLS'
T
tink2123 已提交
663
    ]
T
tink2123 已提交
664
    if alg in ['Rosetta', 'CRNN', 'STARNet', 'RARE', 'SRN']:
L
licx 已提交
665 666 667 668 669 670 671 672
        config['Global']['char_ops'] = CharacterOps(config['Global'])

    place = fluid.CUDAPlace(0) if use_gpu else fluid.CPUPlace()
    startup_program = fluid.Program()
    train_program = fluid.Program()

    if alg in ['EAST', 'DB', 'SAST']:
        train_alg_type = 'det'
W
WenmuZhou 已提交
673
    elif alg in ['Rosetta', 'CRNN', 'STARNet', 'RARE', 'SRN']:
L
licx 已提交
674
        train_alg_type = 'rec'
W
WenmuZhou 已提交
675 676
    else:
        train_alg_type = 'cls'
L
licx 已提交
677 678

    return startup_program, train_program, place, config, train_alg_type