train.py 3.74 KB
Newer Older
LDOUBLEV's avatar
LDOUBLEV committed
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
# 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

import os
import sys
WenmuZhou's avatar
WenmuZhou committed
21

22
__dir__ = os.path.dirname(os.path.abspath(__file__))
LDOUBLEV's avatar
LDOUBLEV committed
23
sys.path.append(__dir__)
24
sys.path.append(os.path.abspath(os.path.join(__dir__, '..')))
LDOUBLEV's avatar
LDOUBLEV committed
25

WenmuZhou's avatar
WenmuZhou committed
26
27
28
import yaml
import paddle
import paddle.distributed as dist
LDOUBLEV's avatar
LDOUBLEV committed
29

WenmuZhou's avatar
WenmuZhou committed
30
paddle.manual_seed(2)
LDOUBLEV's avatar
LDOUBLEV committed
31

WenmuZhou's avatar
WenmuZhou committed
32
33
from ppocr.utils.logging import get_logger
from ppocr.data import build_dataloader
dyning's avatar
dyning committed
34
35
from ppocr.modeling.architectures import build_model
from ppocr.losses import build_loss
WenmuZhou's avatar
WenmuZhou committed
36
37
38
39
40
41
from ppocr.optimizer import build_optimizer
from ppocr.postprocess import build_post_process
from ppocr.metrics import build_metric
from ppocr.utils.save_load import init_model
from ppocr.utils.utility import print_dict
import tools.program as program
LDOUBLEV's avatar
LDOUBLEV committed
42

WenmuZhou's avatar
WenmuZhou committed
43
dist.get_world_size()
LDOUBLEV's avatar
LDOUBLEV committed
44
45


WenmuZhou's avatar
WenmuZhou committed
46
47
48
49
def main(config, device, logger, vdl_writer):
    # init dist environment
    if config['Global']['distributed']:
        dist.init_parallel_env()
LDOUBLEV's avatar
LDOUBLEV committed
50

WenmuZhou's avatar
WenmuZhou committed
51
    global_config = config['Global']
dyning's avatar
dyning committed
52
    
WenmuZhou's avatar
WenmuZhou committed
53
    # build dataloader
dyning's avatar
dyning committed
54
55
56
    train_dataloader = build_dataloader(config, 'Train', device)
    if config['Eval']:
        valid_dataloader = build_dataloader(config, 'Eval', device)
WenmuZhou's avatar
WenmuZhou committed
57
    else:
dyning's avatar
dyning committed
58
59
        valid_dataloader = None

WenmuZhou's avatar
WenmuZhou committed
60
    # build post process
dyning's avatar
dyning committed
61
62
63
    post_process_class = build_post_process(
        config['PostProcess'], global_config)
    
WenmuZhou's avatar
WenmuZhou committed
64
    # build model
dyning's avatar
dyning committed
65
    #for rec algorithm
WenmuZhou's avatar
WenmuZhou committed
66
    if hasattr(post_process_class, 'character'):
dyning's avatar
dyning committed
67
68
        char_num = len(getattr(post_process_class, 'character'))
        config['Architecture']["Head"]['out_channels'] = char_num
WenmuZhou's avatar
WenmuZhou committed
69
70
71
72
    model = build_model(config['Architecture'])
    if config['Global']['distributed']:
        model = paddle.DataParallel(model)

dyning's avatar
dyning committed
73
74
75
    # build loss
    loss_class = build_loss(config['Loss'])
    
WenmuZhou's avatar
WenmuZhou committed
76
    # build optim
dyning's avatar
dyning committed
77
    optimizer, lr_scheduler = build_optimizer(config['Optimizer'],
WenmuZhou's avatar
WenmuZhou committed
78
        epochs=config['Global']['epoch_num'],
dyning's avatar
dyning committed
79
        step_each_epoch=len(train_dataloader),
WenmuZhou's avatar
WenmuZhou committed
80
81
82
83
        parameters=model.parameters())

    # build metric
    eval_class = build_metric(config['Metric'])
dyning's avatar
dyning committed
84
85
86
    
    # load pretrain model
    pre_best_model_dict = init_model(config, model, logger, optimizer)
WenmuZhou's avatar
WenmuZhou committed
87
88

    # start train
dyning's avatar
dyning committed
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
    program.train(config,
        train_dataloader,
        valid_dataloader,
        device,
        model,
        loss_class,
        optimizer,
        lr_scheduler,
        post_process_class,
        eval_class,
        pre_best_model_dict,
        logger,
        vdl_writer)


def test_reader(config, device, logger):
    loader = build_dataloader(config, 'Train', device)
#     loader = build_dataloader(config, 'Eval', device)
107
108
109
110
    import time
    starttime = time.time()
    count = 0
    try:
dyning's avatar
dyning committed
111
        for data in loader():
112
113
114
115
            count += 1
            if count % 1 == 0:
                batch_time = time.time() - starttime
                starttime = time.time()
dyning's avatar
dyning committed
116
                logger.info("reader: {}, {}, {}".format(count, len(data), batch_time))
117
    except Exception as e:
LDOUBLEV's avatar
LDOUBLEV committed
118
119
        logger.info(e)
    logger.info("finish reader: {}, Success!".format(count))
120

LDOUBLEV's avatar
LDOUBLEV committed
121
if __name__ == '__main__':
dyning's avatar
dyning committed
122
123
124
    config, device, logger, vdl_writer = program.preprocess()
    main(config, device, logger, vdl_writer)
#     test_reader(config, device, logger)