program.py 19.1 KB
Newer Older
MissPenguin's avatar
refine  
MissPenguin committed
1
# Copyright (c) 2021 PaddlePaddle Authors. All Rights Reserved.
LDOUBLEV's avatar
LDOUBLEV committed
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
#
# 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

WenmuZhou's avatar
WenmuZhou committed
19
import os
LDOUBLEV's avatar
LDOUBLEV committed
20
import sys
21
import platform
LDOUBLEV's avatar
LDOUBLEV committed
22
23
import yaml
import time
WenmuZhou's avatar
WenmuZhou committed
24
25
26
27
28
29
import shutil
import paddle
import paddle.distributed as dist
from tqdm import tqdm
from argparse import ArgumentParser, RawDescriptionHelpFormatter

LDOUBLEV's avatar
LDOUBLEV committed
30
31
from ppocr.utils.stats import TrainingStats
from ppocr.utils.save_load import save_model
dyning's avatar
dyning committed
32
33
from ppocr.utils.utility import print_dict
from ppocr.utils.logging import get_logger
LDOUBLEV's avatar
LDOUBLEV committed
34
from ppocr.utils import profiler
dyning's avatar
dyning committed
35
36
from ppocr.data import build_dataloader
import numpy as np
LDOUBLEV's avatar
LDOUBLEV committed
37

dyning's avatar
dyning committed
38

LDOUBLEV's avatar
LDOUBLEV committed
39
40
41
42
43
44
45
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")
LDOUBLEV's avatar
LDOUBLEV committed
46
47
48
        self.add_argument(
            '-p',
            '--profiler_options',
49
50
            type=bool,
            default=False,
LDOUBLEV's avatar
LDOUBLEV committed
51
52
            help='The option of profiler, which should be in format \"key1=value1;key2=value2;key3=value3\".'
        )
LDOUBLEV's avatar
LDOUBLEV committed
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
80

    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


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
    """
    _, ext = os.path.splitext(file_path)
    assert ext in ['.yml', '.yaml'], "only support yaml files for now"
81
82
    config = yaml.load(open(file_path, 'rb'), Loader=yaml.Loader)
    return config
LDOUBLEV's avatar
LDOUBLEV committed
83
84


85
def merge_config(config, opts):
LDOUBLEV's avatar
LDOUBLEV committed
86
87
88
89
90
91
    """
    Merge config into global config.
    Args:
        config (dict): Config to be merged.
    Returns: global config
    """
92
    for key, value in opts.items():
LDOUBLEV's avatar
LDOUBLEV committed
93
        if "." not in key:
94
95
            if isinstance(value, dict) and key in config:
                config[key].update(value)
LDOUBLEV's avatar
LDOUBLEV committed
96
            else:
97
                config[key] = value
LDOUBLEV's avatar
LDOUBLEV committed
98
99
        else:
            sub_keys = key.split('.')
tink2123's avatar
tink2123 committed
100
            assert (
101
                sub_keys[0] in config
tink2123's avatar
tink2123 committed
102
            ), "the sub_keys can only be one of global_config: {}, but get: {}, please check your running command".format(
103
104
                config.keys(), sub_keys[0])
            cur = config[sub_keys[0]]
LDOUBLEV's avatar
LDOUBLEV committed
105
106
107
108
109
            for idx, sub_key in enumerate(sub_keys[1:]):
                if idx == len(sub_keys) - 2:
                    cur[sub_key] = value
                else:
                    cur = cur[sub_key]
110
    return config
LDOUBLEV's avatar
LDOUBLEV committed
111
112
113
114
115
116
117
118
119
120
121
122
123
124


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:
WenmuZhou's avatar
WenmuZhou committed
125
        if use_gpu and not paddle.is_compiled_with_cuda():
WenmuZhou's avatar
WenmuZhou committed
126
            print(err)
LDOUBLEV's avatar
LDOUBLEV committed
127
128
129
130
131
            sys.exit(1)
    except Exception as e:
        pass


WenmuZhou's avatar
WenmuZhou committed
132
def train(config,
dyning's avatar
dyning committed
133
134
135
          train_dataloader,
          valid_dataloader,
          device,
WenmuZhou's avatar
WenmuZhou committed
136
137
138
139
140
141
142
143
          model,
          loss_class,
          optimizer,
          lr_scheduler,
          post_process_class,
          eval_class,
          pre_best_model_dict,
          logger,
stephon's avatar
stephon committed
144
145
          vdl_writer=None,
          scaler=None):
WenmuZhou's avatar
WenmuZhou committed
146
147
    cal_metric_during_train = config['Global'].get('cal_metric_during_train',
                                                   False)
LDOUBLEV's avatar
LDOUBLEV committed
148
149
150
151
    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']
LDOUBLEV's avatar
LDOUBLEV committed
152
    profiler_options = config['profiler_options']
153
154
155
156
    if profiler_options is True:
        profiler_options = "batch_range=[10,20];state=GPU;tracer_option=Default;profile_path=model.profile"
    else:
        profiler_options = None
WenmuZhou's avatar
WenmuZhou committed
157

dyning's avatar
dyning committed
158
    global_step = 0
159
160
    if 'global_step' in pre_best_model_dict:
        global_step = pre_best_model_dict['global_step']
LDOUBLEV's avatar
LDOUBLEV committed
161
162
163
164
    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]
WenmuZhou's avatar
WenmuZhou committed
165
166
167
168
169
        if len(valid_dataloader) == 0:
            logger.info(
                'No Images in eval dataset, evaluation during training will be disabled'
            )
            start_eval_step = 1e111
LDOUBLEV's avatar
LDOUBLEV committed
170
171
172
        logger.info(
            "During the training process, after the {}th iteration, an evaluation is run every {} iterations".
            format(start_eval_step, eval_batch_step))
LDOUBLEV's avatar
LDOUBLEV committed
173
174
    save_epoch_step = config['Global']['save_epoch_step']
    save_model_dir = config['Global']['save_model_dir']
175
176
    if not os.path.exists(save_model_dir):
        os.makedirs(save_model_dir)
WenmuZhou's avatar
WenmuZhou committed
177
178
179
180
    main_indicator = eval_class.main_indicator
    best_model_dict = {main_indicator: 0}
    best_model_dict.update(pre_best_model_dict)
    train_stats = TrainingStats(log_smooth_window, ['lr'])
tink2123's avatar
tink2123 committed
181
    model_average = False
WenmuZhou's avatar
WenmuZhou committed
182
183
    model.train()

tink2123's avatar
tink2123 committed
184
    use_srn = config['Architecture']['algorithm'] == "SRN"
tink2123's avatar
tink2123 committed
185
    extra_input = config['Architecture'][
LDOUBLEV's avatar
LDOUBLEV committed
186
        'algorithm'] in ["SRN", "NRTR", "SAR", "SEED"]
187
    try:
LDOUBLEV's avatar
fix bug  
LDOUBLEV committed
188
        model_type = config['Architecture']['model_type']
189
    except:
LDOUBLEV's avatar
fix bug  
LDOUBLEV committed
190
        model_type = None
tink2123's avatar
tink2123 committed
191
    algorithm = config['Architecture']['algorithm']
tink2123's avatar
tink2123 committed
192

193
194
195
196
197
198
199
200
201
202
    start_epoch = best_model_dict[
        'start_epoch'] if 'start_epoch' in best_model_dict else 1

    train_reader_cost = 0.0
    train_run_cost = 0.0
    total_samples = 0
    reader_start = time.time()

    max_iter = len(train_dataloader) - 1 if platform.system(
    ) == "Windows" else len(train_dataloader)
WenmuZhou's avatar
WenmuZhou committed
203

tink2123's avatar
tink2123 committed
204
    for epoch in range(start_epoch, epoch_num + 1):
205
206
207
208
209
210
        if train_dataloader.dataset.need_reset:
            train_dataloader = build_dataloader(
                config, 'Train', device, logger, seed=epoch)
            max_iter = len(train_dataloader) - 1 if platform.system(
            ) == "Windows" else len(train_dataloader)

WenmuZhou's avatar
WenmuZhou committed
211
        for idx, batch in enumerate(train_dataloader):
LDOUBLEV's avatar
LDOUBLEV committed
212
            profiler.add_profiler_step(profiler_options)
WenmuZhou's avatar
WenmuZhou committed
213
            train_reader_cost += time.time() - reader_start
Jane-Ding's avatar
Jane-Ding committed
214
            if idx >= max_iter:
WenmuZhou's avatar
WenmuZhou committed
215
216
217
                break
            lr = optimizer.get_lr()
            images = batch[0]
tink2123's avatar
tink2123 committed
218
            if use_srn:
tink2123's avatar
tink2123 committed
219
                model_average = True
stephon's avatar
stephon committed
220

WenmuZhou's avatar
WenmuZhou committed
221
            train_start = time.time()
stephon's avatar
stephon committed
222
223
224
225
226
227
228
            # use amp
            if scaler:
                with paddle.amp.auto_cast():
                    if model_type == 'table' or extra_input:
                        preds = model(images, data=batch[1:])
                    else:
                        preds = model(images)
tink2123's avatar
tink2123 committed
229
            else:
stephon's avatar
stephon committed
230
231
                if model_type == 'table' or extra_input:
                    preds = model(images, data=batch[1:])
232
                elif model_type in ["kie", 'vqa']:
LDOUBLEV's avatar
LDOUBLEV committed
233
                    preds = model(batch)
stephon's avatar
stephon committed
234
235
                else:
                    preds = model(images)
236

WenmuZhou's avatar
WenmuZhou committed
237
238
            loss = loss_class(preds, batch)
            avg_loss = loss['loss']
stephon's avatar
stephon committed
239
240
241
242
243
244
245
246

            if scaler:
                scaled_avg_loss = scaler.scale(avg_loss)
                scaled_avg_loss.backward()
                scaler.minimize(optimizer, scaled_avg_loss)
            else:
                avg_loss.backward()
                optimizer.step()
WenmuZhou's avatar
WenmuZhou committed
247
            optimizer.clear_grad()
WenmuZhou's avatar
WenmuZhou committed
248

WenmuZhou's avatar
WenmuZhou committed
249
            train_run_cost += time.time() - train_start
250
            global_step += 1
WenmuZhou's avatar
WenmuZhou committed
251
            total_samples += len(images)
WenmuZhou's avatar
WenmuZhou committed
252

dyning's avatar
dyning committed
253
254
            if not isinstance(lr_scheduler, float):
                lr_scheduler.step()
WenmuZhou's avatar
WenmuZhou committed
255
256
257
258
259
260

            # logger and visualdl
            stats = {k: v.numpy().mean() for k, v in loss.items()}
            stats['lr'] = lr
            train_stats.update(stats)

LDOUBLEV's avatar
LDOUBLEV committed
261
            if cal_metric_during_train:  # only rec and cls need
WenmuZhou's avatar
WenmuZhou committed
262
                batch = [item.numpy() for item in batch]
LDOUBLEV's avatar
LDOUBLEV committed
263
                if model_type in ['table', 'kie']:
MissPenguin's avatar
MissPenguin committed
264
265
266
267
                    eval_class(preds, batch)
                else:
                    post_result = post_process_class(preds, batch[1])
                    eval_class(post_result, batch)
littletomatodonkey's avatar
fix doc  
littletomatodonkey committed
268
269
                metric = eval_class.get_metric()
                train_stats.update(metric)
WenmuZhou's avatar
WenmuZhou committed
270
271
272
273
274
275

            if vdl_writer is not None and dist.get_rank() == 0:
                for k, v in train_stats.get().items():
                    vdl_writer.add_scalar('TRAIN/{}'.format(k), v, global_step)
                vdl_writer.add_scalar('TRAIN/lr', lr, global_step)

276
277
278
            if dist.get_rank() == 0 and (
                (global_step > 0 and global_step % print_batch_step == 0) or
                (idx >= len(train_dataloader) - 1)):
WenmuZhou's avatar
WenmuZhou committed
279
                logs = train_stats.log()
280
                strs = 'epoch: [{}/{}], global_step: {}, {}, avg_reader_cost: {:.5f} s, avg_batch_cost: {:.5f} s, avg_samples: {}, ips: {:.5f}'.format(
WenmuZhou's avatar
WenmuZhou committed
281
                    epoch, epoch_num, global_step, logs, train_reader_cost /
WenmuZhou's avatar
WenmuZhou committed
282
                    print_batch_step, (train_reader_cost + train_run_cost) /
283
                    print_batch_step, total_samples / print_batch_step,
WenmuZhou's avatar
WenmuZhou committed
284
                    total_samples / (train_reader_cost + train_run_cost))
WenmuZhou's avatar
WenmuZhou committed
285
                logger.info(strs)
286

WenmuZhou's avatar
WenmuZhou committed
287
                train_reader_cost = 0.0
WenmuZhou's avatar
WenmuZhou committed
288
289
                train_run_cost = 0.0
                total_samples = 0
WenmuZhou's avatar
WenmuZhou committed
290
291
292
            # eval
            if global_step > start_eval_step and \
                    (global_step - start_eval_step) % eval_batch_step == 0 and dist.get_rank() == 0:
tink2123's avatar
tink2123 committed
293
294
295
296
297
298
299
                if model_average:
                    Model_Average = paddle.incubate.optimizer.ModelAverage(
                        0.15,
                        parameters=model.parameters(),
                        min_average_window=10000,
                        max_average_window=15625)
                    Model_Average.apply()
tink2123's avatar
tink2123 committed
300
301
302
303
304
                cur_metric = eval(
                    model,
                    valid_dataloader,
                    post_process_class,
                    eval_class,
MissPenguin's avatar
refine  
MissPenguin committed
305
                    model_type,
tink2123's avatar
tink2123 committed
306
                    extra_input=extra_input)
LDOUBLEV's avatar
LDOUBLEV committed
307
308
309
                cur_metric_str = 'cur metric, {}'.format(', '.join(
                    ['{}: {}'.format(k, v) for k, v in cur_metric.items()]))
                logger.info(cur_metric_str)
WenmuZhou's avatar
WenmuZhou committed
310
311
312

                # logger metric
                if vdl_writer is not None:
LDOUBLEV's avatar
LDOUBLEV committed
313
                    for k, v in cur_metric.items():
WenmuZhou's avatar
WenmuZhou committed
314
315
                        if isinstance(v, (float, int)):
                            vdl_writer.add_scalar('EVAL/{}'.format(k),
LDOUBLEV's avatar
LDOUBLEV committed
316
317
                                                  cur_metric[k], global_step)
                if cur_metric[main_indicator] >= best_model_dict[
WenmuZhou's avatar
WenmuZhou committed
318
                        main_indicator]:
LDOUBLEV's avatar
LDOUBLEV committed
319
                    best_model_dict.update(cur_metric)
WenmuZhou's avatar
WenmuZhou committed
320
321
322
323
324
325
                    best_model_dict['best_epoch'] = epoch
                    save_model(
                        model,
                        optimizer,
                        save_model_dir,
                        logger,
326
                        config,
WenmuZhou's avatar
WenmuZhou committed
327
328
329
                        is_best=True,
                        prefix='best_accuracy',
                        best_model_dict=best_model_dict,
330
331
                        epoch=epoch,
                        global_step=global_step)
LDOUBLEV's avatar
LDOUBLEV committed
332
                best_str = 'best metric, {}'.format(', '.join([
WenmuZhou's avatar
WenmuZhou committed
333
334
335
336
337
338
339
340
                    '{}: {}'.format(k, v) for k, v in best_model_dict.items()
                ]))
                logger.info(best_str)
                # logger best metric
                if vdl_writer is not None:
                    vdl_writer.add_scalar('EVAL/best_{}'.format(main_indicator),
                                          best_model_dict[main_indicator],
                                          global_step)
341

WenmuZhou's avatar
WenmuZhou committed
342
            reader_start = time.time()
WenmuZhou's avatar
WenmuZhou committed
343
344
345
346
347
348
        if dist.get_rank() == 0:
            save_model(
                model,
                optimizer,
                save_model_dir,
                logger,
349
                config,
WenmuZhou's avatar
WenmuZhou committed
350
351
352
                is_best=False,
                prefix='latest',
                best_model_dict=best_model_dict,
353
354
                epoch=epoch,
                global_step=global_step)
WenmuZhou's avatar
WenmuZhou committed
355
356
357
358
359
360
        if dist.get_rank() == 0 and epoch > 0 and epoch % save_epoch_step == 0:
            save_model(
                model,
                optimizer,
                save_model_dir,
                logger,
361
                config,
WenmuZhou's avatar
WenmuZhou committed
362
363
364
                is_best=False,
                prefix='iter_epoch_{}'.format(epoch),
                best_model_dict=best_model_dict,
365
366
                epoch=epoch,
                global_step=global_step)
LDOUBLEV's avatar
LDOUBLEV committed
367
    best_str = 'best metric, {}'.format(', '.join(
WenmuZhou's avatar
WenmuZhou committed
368
369
370
371
        ['{}: {}'.format(k, v) for k, v in best_model_dict.items()]))
    logger.info(best_str)
    if dist.get_rank() == 0 and vdl_writer is not None:
        vdl_writer.close()
LDOUBLEV's avatar
LDOUBLEV committed
372
373
374
    return


MissPenguin's avatar
refine  
MissPenguin committed
375
376
377
378
def eval(model,
         valid_dataloader,
         post_process_class,
         eval_class,
LDOUBLEV's avatar
LDOUBLEV committed
379
         model_type=None,
tink2123's avatar
tink2123 committed
380
         extra_input=False):
WenmuZhou's avatar
WenmuZhou committed
381
382
383
384
    model.eval()
    with paddle.no_grad():
        total_frame = 0.0
        total_time = 0.0
WenmuZhou's avatar
WenmuZhou committed
385
386
387
388
389
        pbar = tqdm(
            total=len(valid_dataloader),
            desc='eval model:',
            position=0,
            leave=True)
390
391
        max_iter = len(valid_dataloader) - 1 if platform.system(
        ) == "Windows" else len(valid_dataloader)
WenmuZhou's avatar
WenmuZhou committed
392
        for idx, batch in enumerate(valid_dataloader):
393
            if idx >= max_iter:
WenmuZhou's avatar
WenmuZhou committed
394
                break
WenmuZhou's avatar
fix bug  
WenmuZhou committed
395
            images = batch[0]
WenmuZhou's avatar
WenmuZhou committed
396
            start = time.time()
tink2123's avatar
tink2123 committed
397
            if model_type == 'table' or extra_input:
MissPenguin's avatar
refine  
MissPenguin committed
398
                preds = model(images, data=batch[1:])
399
            elif model_type in ["kie", 'vqa']:
LDOUBLEV's avatar
LDOUBLEV committed
400
                preds = model(batch)
xiaoting's avatar
xiaoting committed
401
            else:
LDOUBLEV's avatar
LDOUBLEV committed
402
                preds = model(images)
403
404
405
406
407
408
409

            batch_numpy = []
            for item in batch:
                if isinstance(item, paddle.Tensor):
                    batch_numpy.append(item.numpy())
                else:
                    batch_numpy.append(item)
WenmuZhou's avatar
WenmuZhou committed
410
411
412
            # Obtain usable results from post-processing methods
            total_time += time.time() - start
            # Evaluate the results of the current batch
LDOUBLEV's avatar
LDOUBLEV committed
413
            if model_type in ['table', 'kie']:
414
415
416
417
                eval_class(preds, batch_numpy)
            elif model_type in ['vqa']:
                post_result = post_process_class(preds, batch_numpy)
                eval_class(post_result, batch_numpy)
MissPenguin's avatar
MissPenguin committed
418
            else:
419
420
                post_result = post_process_class(preds, batch_numpy[1])
                eval_class(post_result, batch_numpy)
LDOUBLEV's avatar
LDOUBLEV committed
421

WenmuZhou's avatar
fix bug  
WenmuZhou committed
422
            pbar.update(1)
WenmuZhou's avatar
WenmuZhou committed
423
            total_frame += len(images)
LDOUBLEV's avatar
LDOUBLEV committed
424
425
        # Get final metric,eg. acc or hmean
        metric = eval_class.get_metric()
dyning's avatar
dyning committed
426

WenmuZhou's avatar
fix bug  
WenmuZhou committed
427
    pbar.close()
WenmuZhou's avatar
WenmuZhou committed
428
    model.train()
LDOUBLEV's avatar
LDOUBLEV committed
429
430
    metric['fps'] = total_frame / total_time
    return metric
licx's avatar
licx committed
431

tink2123's avatar
tink2123 committed
432

Bin Lu's avatar
Bin Lu committed
433
434
435
436
437
438
439
440
441
442
443
444
445
446
447
448
449
450
451
452
453
454
455
456
457
458
459
460
461
462
463
464
465
466
467
468
469
470
471
472
473
474
475
476
477
478
479
480
481
def update_center(char_center, post_result, preds):
    result, label = post_result
    feats, logits = preds
    logits = paddle.argmax(logits, axis=-1)
    feats = feats.numpy()
    logits = logits.numpy()

    for idx_sample in range(len(label)):
        if result[idx_sample][0] == label[idx_sample][0]:
            feat = feats[idx_sample]
            logit = logits[idx_sample]
            for idx_time in range(len(logit)):
                index = logit[idx_time]
                if index in char_center.keys():
                    char_center[index][0] = (
                        char_center[index][0] * char_center[index][1] +
                        feat[idx_time]) / (char_center[index][1] + 1)
                    char_center[index][1] += 1
                else:
                    char_center[index] = [feat[idx_time], 1]
    return char_center


def get_center(model, eval_dataloader, post_process_class):
    pbar = tqdm(total=len(eval_dataloader), desc='get center:')
    max_iter = len(eval_dataloader) - 1 if platform.system(
    ) == "Windows" else len(eval_dataloader)
    char_center = dict()
    for idx, batch in enumerate(eval_dataloader):
        if idx >= max_iter:
            break
        images = batch[0]
        start = time.time()
        preds = model(images)

        batch = [item.numpy() for item in batch]
        # Obtain usable results from post-processing methods
        post_result = post_process_class(preds, batch[1])

        #update char_center
        char_center = update_center(char_center, post_result, preds)
        pbar.update(1)

    pbar.close()
    for key in char_center.keys():
        char_center[key] = char_center[key][0]
    return char_center


482
def preprocess(is_train=False):
licx's avatar
licx committed
483
    FLAGS = ArgsParser().parse_args()
LDOUBLEV's avatar
LDOUBLEV committed
484
    profiler_options = FLAGS.profiler_options
licx's avatar
licx committed
485
    config = load_config(FLAGS.config)
486
    config = merge_config(config, FLAGS.opt)
LDOUBLEV's avatar
LDOUBLEV committed
487
    profile_dic = {"profiler_options": FLAGS.profiler_options}
488
    config = merge_config(config, profile_dic)
licx's avatar
licx committed
489

490
491
492
493
494
495
496
497
498
499
500
    if is_train:
        # save_config
        save_model_dir = config['Global']['save_model_dir']
        os.makedirs(save_model_dir, exist_ok=True)
        with open(os.path.join(save_model_dir, 'config.yml'), 'w') as f:
            yaml.dump(
                dict(config), f, default_flow_style=False, sort_keys=False)
        log_file = '{}/train.log'.format(save_model_dir)
    else:
        log_file = None
    logger = get_logger(name='root', log_file=log_file)
licx's avatar
licx committed
501
502
503
504
505

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

WenmuZhou's avatar
WenmuZhou committed
506
507
    alg = config['Architecture']['algorithm']
    assert alg in [
Jethong's avatar
Jethong committed
508
        'EAST', 'DB', 'SAST', 'Rosetta', 'CRNN', 'STARNet', 'RARE', 'SRN',
tink2123's avatar
tink2123 committed
509
        'CLS', 'PGNet', 'Distillation', 'NRTR', 'TableAttn', 'SAR', 'PSE',
510
        'SEED', 'SDMGR', 'LayoutXLM', 'LayoutLM'
WenmuZhou's avatar
WenmuZhou committed
511
    ]
licx's avatar
licx committed
512

WenmuZhou's avatar
WenmuZhou committed
513
514
    device = 'gpu:{}'.format(dist.ParallelEnv().dev_id) if use_gpu else 'cpu'
    device = paddle.set_device(device)
dyning's avatar
dyning committed
515

dyning's avatar
dyning committed
516
    config['Global']['distributed'] = dist.get_world_size() != 1
517

dyning's avatar
dyning committed
518
519
    if config['Global']['use_visualdl']:
        from visualdl import LogWriter
LDOUBLEV's avatar
fix bug  
LDOUBLEV committed
520
        save_model_dir = config['Global']['save_model_dir']
dyning's avatar
dyning committed
521
522
523
524
525
526
527
528
529
        vdl_writer_path = '{}/vdl/'.format(save_model_dir)
        os.makedirs(vdl_writer_path, exist_ok=True)
        vdl_writer = LogWriter(logdir=vdl_writer_path)
    else:
        vdl_writer = None
    print_dict(config, logger)
    logger.info('train with paddle {} and device {}'.format(paddle.__version__,
                                                            device))
    return config, device, logger, vdl_writer