program.py 18.7 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
49
50
51
52
        self.add_argument(
            '-p',
            '--profiler_options',
            type=str,
            default=None,
            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
81
82
83
84
85
86

    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()

lyl120117's avatar
lyl120117 committed
87
88
default_config = {'Global': {'debug': False, }}

LDOUBLEV's avatar
LDOUBLEV committed
89
90
91
92
93
94
95
96

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
    """
lyl120117's avatar
lyl120117 committed
97
    merge_config(default_config)
LDOUBLEV's avatar
LDOUBLEV committed
98
99
    _, ext = os.path.splitext(file_path)
    assert ext in ['.yml', '.yaml'], "only support yaml files for now"
WenmuZhou's avatar
WenmuZhou committed
100
    merge_config(yaml.load(open(file_path, 'rb'), Loader=yaml.Loader))
LDOUBLEV's avatar
LDOUBLEV committed
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
    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('.')
tink2123's avatar
tink2123 committed
119
120
121
122
            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])
LDOUBLEV's avatar
LDOUBLEV committed
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
            cur = global_config[sub_keys[0]]
            for idx, sub_key in enumerate(sub_keys[1:]):
                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:
WenmuZhou's avatar
WenmuZhou committed
143
        if use_gpu and not paddle.is_compiled_with_cuda():
WenmuZhou's avatar
WenmuZhou committed
144
            print(err)
LDOUBLEV's avatar
LDOUBLEV committed
145
146
147
148
149
            sys.exit(1)
    except Exception as e:
        pass


WenmuZhou's avatar
WenmuZhou committed
150
def train(config,
dyning's avatar
dyning committed
151
152
153
          train_dataloader,
          valid_dataloader,
          device,
WenmuZhou's avatar
WenmuZhou committed
154
155
156
157
158
159
160
161
162
163
164
          model,
          loss_class,
          optimizer,
          lr_scheduler,
          post_process_class,
          eval_class,
          pre_best_model_dict,
          logger,
          vdl_writer=None):
    cal_metric_during_train = config['Global'].get('cal_metric_during_train',
                                                   False)
LDOUBLEV's avatar
LDOUBLEV committed
165
166
167
168
    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
169
    profiler_options = config['profiler_options']
WenmuZhou's avatar
WenmuZhou committed
170

dyning's avatar
dyning committed
171
    global_step = 0
172
173
    if 'global_step' in pre_best_model_dict:
        global_step = pre_best_model_dict['global_step']
LDOUBLEV's avatar
LDOUBLEV committed
174
175
176
177
    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
178
179
180
181
182
        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
183
184
185
        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
186
187
    save_epoch_step = config['Global']['save_epoch_step']
    save_model_dir = config['Global']['save_model_dir']
188
189
    if not os.path.exists(save_model_dir):
        os.makedirs(save_model_dir)
WenmuZhou's avatar
WenmuZhou committed
190
191
192
193
    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
194
    model_average = False
WenmuZhou's avatar
WenmuZhou committed
195
196
    model.train()

tink2123's avatar
tink2123 committed
197
    use_srn = config['Architecture']['algorithm'] == "SRN"
tink2123's avatar
tink2123 committed
198
199
    extra_input = config['Architecture'][
        'algorithm'] in ["SRN", "NRTR", "SAR", "SEED"]
200
    try:
LDOUBLEV's avatar
fix bug  
LDOUBLEV committed
201
        model_type = config['Architecture']['model_type']
202
    except:
LDOUBLEV's avatar
fix bug  
LDOUBLEV committed
203
        model_type = None
tink2123's avatar
tink2123 committed
204
    algorithm = config['Architecture']['algorithm']
MissPenguin's avatar
refine  
MissPenguin committed
205

WenmuZhou's avatar
WenmuZhou committed
206
207
208
    if 'start_epoch' in best_model_dict:
        start_epoch = best_model_dict['start_epoch']
    else:
tink2123's avatar
tink2123 committed
209
        start_epoch = 1
WenmuZhou's avatar
WenmuZhou committed
210

tink2123's avatar
tink2123 committed
211
    for epoch in range(start_epoch, epoch_num + 1):
212
213
        train_dataloader = build_dataloader(
            config, 'Train', device, logger, seed=epoch)
WenmuZhou's avatar
WenmuZhou committed
214
215
216
217
        train_batch_cost = 0.0
        train_reader_cost = 0.0
        batch_sum = 0
        batch_start = time.time()
Jane-Ding's avatar
Jane-Ding committed
218
219
220
        max_iter = len(train_dataloader) - 1 if platform.system(
        ) == "Windows" else len(train_dataloader)
        for idx, batch in enumerate(train_dataloader):
LDOUBLEV's avatar
LDOUBLEV committed
221
            profiler.add_profiler_step(profiler_options)
WenmuZhou's avatar
WenmuZhou committed
222
            train_reader_cost += time.time() - batch_start
Jane-Ding's avatar
Jane-Ding committed
223
224
            if idx >= max_iter:
                break
WenmuZhou's avatar
WenmuZhou committed
225
226
            lr = optimizer.get_lr()
            images = batch[0]
tink2123's avatar
tink2123 committed
227
            if use_srn:
tink2123's avatar
tink2123 committed
228
                model_average = True
stephon's avatar
stephon committed
229
230
231
232
233
234
235
236

            # 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)
MissPenguin's avatar
refine  
MissPenguin committed
237
            else:
stephon's avatar
stephon committed
238
239
240
241
                if model_type == 'table' or extra_input:
                    preds = model(images, data=batch[1:])
                else:
                    preds = model(images)
WenmuZhou's avatar
WenmuZhou committed
242
243
            loss = loss_class(preds, batch)
            avg_loss = loss['loss']
stephon's avatar
stephon committed
244
245
246
247
248
249
250
251

            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
252
            optimizer.clear_grad()
WenmuZhou's avatar
WenmuZhou committed
253
254
255
256

            train_batch_cost += time.time() - batch_start
            batch_sum += len(images)

dyning's avatar
dyning committed
257
258
            if not isinstance(lr_scheduler, float):
                lr_scheduler.step()
WenmuZhou's avatar
WenmuZhou committed
259
260
261
262
263
264

            # 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
265
            if cal_metric_during_train:  # only rec and cls need
WenmuZhou's avatar
WenmuZhou committed
266
                batch = [item.numpy() for item in batch]
MissPenguin's avatar
MissPenguin committed
267
268
269
270
271
                if model_type == 'table':
                    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
272
273
                metric = eval_class.get_metric()
                train_stats.update(metric)
WenmuZhou's avatar
WenmuZhou committed
274
275
276
277
278
279

            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)

280
281
282
            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
283
                logs = train_stats.log()
WenmuZhou's avatar
WenmuZhou committed
284
                strs = 'epoch: [{}/{}], iter: {}, {}, reader_cost: {:.5f} s, batch_cost: {:.5f} s, samples: {}, ips: {:.5f}'.format(
WenmuZhou's avatar
WenmuZhou committed
285
286
287
                    epoch, epoch_num, global_step, logs, train_reader_cost /
                    print_batch_step, train_batch_cost / print_batch_step,
                    batch_sum, batch_sum / train_batch_cost)
WenmuZhou's avatar
WenmuZhou committed
288
                logger.info(strs)
WenmuZhou's avatar
WenmuZhou committed
289
290
291
                train_batch_cost = 0.0
                train_reader_cost = 0.0
                batch_sum = 0
WenmuZhou's avatar
WenmuZhou committed
292
293
294
            # 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
295
296
297
298
299
300
301
                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
302
303
304
305
306
                cur_metric = eval(
                    model,
                    valid_dataloader,
                    post_process_class,
                    eval_class,
MissPenguin's avatar
refine  
MissPenguin committed
307
                    model_type,
tink2123's avatar
tink2123 committed
308
                    extra_input=extra_input)
LDOUBLEV's avatar
LDOUBLEV committed
309
310
311
                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
312
313
314

                # logger metric
                if vdl_writer is not None:
LDOUBLEV's avatar
LDOUBLEV committed
315
                    for k, v in cur_metric.items():
WenmuZhou's avatar
WenmuZhou committed
316
317
                        if isinstance(v, (float, int)):
                            vdl_writer.add_scalar('EVAL/{}'.format(k),
LDOUBLEV's avatar
LDOUBLEV committed
318
319
                                                  cur_metric[k], global_step)
                if cur_metric[main_indicator] >= best_model_dict[
WenmuZhou's avatar
WenmuZhou committed
320
                        main_indicator]:
LDOUBLEV's avatar
LDOUBLEV committed
321
                    best_model_dict.update(cur_metric)
WenmuZhou's avatar
WenmuZhou committed
322
323
324
325
326
327
328
329
330
                    best_model_dict['best_epoch'] = epoch
                    save_model(
                        model,
                        optimizer,
                        save_model_dir,
                        logger,
                        is_best=True,
                        prefix='best_accuracy',
                        best_model_dict=best_model_dict,
331
332
                        epoch=epoch,
                        global_step=global_step)
LDOUBLEV's avatar
LDOUBLEV committed
333
                best_str = 'best metric, {}'.format(', '.join([
WenmuZhou's avatar
WenmuZhou committed
334
335
336
337
338
339
340
341
342
                    '{}: {}'.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)
            global_step += 1
tink2123's avatar
tink2123 committed
343
            optimizer.clear_grad()
344
            batch_start = time.time()
WenmuZhou's avatar
WenmuZhou committed
345
346
347
348
349
350
351
352
353
        if dist.get_rank() == 0:
            save_model(
                model,
                optimizer,
                save_model_dir,
                logger,
                is_best=False,
                prefix='latest',
                best_model_dict=best_model_dict,
354
355
                epoch=epoch,
                global_step=global_step)
WenmuZhou's avatar
WenmuZhou committed
356
357
358
359
360
361
362
363
364
        if dist.get_rank() == 0 and epoch > 0 and epoch % save_epoch_step == 0:
            save_model(
                model,
                optimizer,
                save_model_dir,
                logger,
                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
fix bug  
WenmuZhou committed
385
        pbar = tqdm(total=len(valid_dataloader), desc='eval model:')
386
387
        max_iter = len(valid_dataloader) - 1 if platform.system(
        ) == "Windows" else len(valid_dataloader)
WenmuZhou's avatar
WenmuZhou committed
388
        for idx, batch in enumerate(valid_dataloader):
389
            if idx >= max_iter:
WenmuZhou's avatar
WenmuZhou committed
390
                break
WenmuZhou's avatar
fix bug  
WenmuZhou committed
391
            images = batch[0]
WenmuZhou's avatar
WenmuZhou committed
392
            start = time.time()
tink2123's avatar
tink2123 committed
393
            if model_type == 'table' or extra_input:
MissPenguin's avatar
refine  
MissPenguin committed
394
395
396
                preds = model(images, data=batch[1:])
            else:
                preds = model(images)
WenmuZhou's avatar
WenmuZhou committed
397
398
399
400
            batch = [item.numpy() for item in batch]
            # Obtain usable results from post-processing methods
            total_time += time.time() - start
            # Evaluate the results of the current batch
MissPenguin's avatar
MissPenguin committed
401
402
403
404
405
            if model_type == 'table':
                eval_class(preds, batch)
            else:
                post_result = post_process_class(preds, batch[1])
                eval_class(post_result, batch)
WenmuZhou's avatar
fix bug  
WenmuZhou committed
406
            pbar.update(1)
WenmuZhou's avatar
WenmuZhou committed
407
            total_frame += len(images)
LDOUBLEV's avatar
LDOUBLEV committed
408
409
        # Get final metric,eg. acc or hmean
        metric = eval_class.get_metric()
dyning's avatar
dyning committed
410

WenmuZhou's avatar
fix bug  
WenmuZhou committed
411
    pbar.close()
WenmuZhou's avatar
WenmuZhou committed
412
    model.train()
LDOUBLEV's avatar
LDOUBLEV committed
413
414
    metric['fps'] = total_frame / total_time
    return metric
licx's avatar
licx committed
415

tink2123's avatar
tink2123 committed
416

Bin Lu's avatar
Bin Lu committed
417
418
419
420
421
422
423
424
425
426
427
428
429
430
431
432
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
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
        total_time += time.time() - start
        # Evaluate the results of the current batch
        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


468
def preprocess(is_train=False):
licx's avatar
licx committed
469
    FLAGS = ArgsParser().parse_args()
LDOUBLEV's avatar
LDOUBLEV committed
470
    profiler_options = FLAGS.profiler_options
licx's avatar
licx committed
471
472
    config = load_config(FLAGS.config)
    merge_config(FLAGS.opt)
LDOUBLEV's avatar
LDOUBLEV committed
473
474
    profile_dic = {"profiler_options": FLAGS.profiler_options}
    merge_config(profile_dic)
licx's avatar
licx committed
475

476
477
478
479
480
481
482
483
484
485
486
487
    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
488
489
490
491
    # 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
492
493
    alg = config['Architecture']['algorithm']
    assert alg in [
Jethong's avatar
Jethong committed
494
        'EAST', 'DB', 'SAST', 'Rosetta', 'CRNN', 'STARNet', 'RARE', 'SRN',
tink2123's avatar
tink2123 committed
495
        'CLS', 'PGNet', 'Distillation', 'NRTR', 'TableAttn', 'SAR', 'PSE',
Bin Lu's avatar
Bin Lu committed
496
497
        'SEED'
    ]
498
499
500
501
502
    windows_not_support_list = ['PSE']
    if platform.system() == "Windows" and alg in windows_not_support_list:
        logger.warning('{} is not support in Windows now'.format(
            windows_not_support_list))
        sys.exit()
licx's avatar
licx committed
503

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

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

dyning's avatar
dyning committed
509
510
    if config['Global']['use_visualdl']:
        from visualdl import LogWriter
LDOUBLEV's avatar
fix bug  
LDOUBLEV committed
511
        save_model_dir = config['Global']['save_model_dir']
dyning's avatar
dyning committed
512
513
514
515
516
517
518
519
520
        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