PaddleOCR/tools/train.py

138 lines
4.6 KiB
Python
Raw Normal View History

2020-05-10 16:26:57 +08:00
# 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
__dir__ = os.path.dirname(os.path.abspath(__file__))
2020-06-12 13:49:24 +08:00
sys.path.append(__dir__)
sys.path.append(os.path.abspath(os.path.join(__dir__, '..')))
2020-05-10 16:26:57 +08:00
def set_paddle_flags(**kwargs):
for key, value in kwargs.items():
if os.environ.get(key, None) is None:
os.environ[key] = str(value)
# NOTE(paddle-dev): All of these flags should be
# set before `import paddle`. Otherwise, it would
2020-05-19 17:30:22 +08:00
# not take any effect.
2020-05-10 16:26:57 +08:00
set_paddle_flags(
FLAGS_eager_delete_tensor_gb=0, # enable GC to save memory
)
import tools.program as program
2020-05-10 16:26:57 +08:00
from paddle import fluid
from ppocr.utils.utility import initial_logger
2020-11-11 18:41:22 +08:00
from ppocr.utils.utility import enable_static_mode
2020-05-10 16:26:57 +08:00
logger = initial_logger()
from ppocr.data.reader_main import reader_main
from ppocr.utils.save_load import init_model
2020-07-11 12:14:05 +08:00
from paddle.fluid.contrib.model_stat import summary
2020-05-10 16:26:57 +08:00
def main():
# build train program
2020-05-10 16:26:57 +08:00
train_build_outputs = program.build(
config, train_program, startup_program, mode='train')
train_loader = train_build_outputs[0]
train_fetch_name_list = train_build_outputs[1]
train_fetch_varname_list = train_build_outputs[2]
train_opt_loss_name = train_build_outputs[3]
2020-08-14 16:31:13 +08:00
model_average = train_build_outputs[-1]
2020-05-10 16:26:57 +08:00
# build eval program
2020-05-10 16:26:57 +08:00
eval_program = fluid.Program()
eval_build_outputs = program.build(
config, eval_program, startup_program, mode='eval')
eval_fetch_name_list = eval_build_outputs[1]
eval_fetch_varname_list = eval_build_outputs[2]
eval_program = eval_program.clone(for_test=True)
# initialize train reader
2020-05-10 16:26:57 +08:00
train_reader = reader_main(config=config, mode="train")
train_loader.set_sample_list_generator(train_reader, places=place)
# initialize eval reader
2020-05-10 16:26:57 +08:00
eval_reader = reader_main(config=config, mode="eval")
exe = fluid.Executor(place)
exe.run(startup_program)
# compile program for multi-devices
train_compile_program = program.create_multi_devices_program(
train_program, train_opt_loss_name)
2020-07-11 12:14:05 +08:00
# dump mode structure
if config['Global']['debug']:
2020-09-01 13:44:51 +08:00
if train_alg_type == 'rec' and 'attention' in config['Global'][
'loss_type']:
2020-07-11 12:14:05 +08:00
logger.warning('Does not suport dump attention...')
else:
summary(train_program)
2020-05-10 16:26:57 +08:00
init_model(config, train_program, exe)
train_info_dict = {'compile_program':train_compile_program,\
'train_program':train_program,\
'reader':train_loader,\
'fetch_name_list':train_fetch_name_list,\
2020-08-14 16:31:13 +08:00
'fetch_varname_list':train_fetch_varname_list,\
'model_average': model_average}
2020-05-10 16:26:57 +08:00
eval_info_dict = {'program':eval_program,\
'reader':eval_reader,\
'fetch_name_list':eval_fetch_name_list,\
'fetch_varname_list':eval_fetch_varname_list}
2020-07-23 14:04:24 +08:00
if train_alg_type == 'det':
2020-05-10 16:26:57 +08:00
program.train_eval_det_run(config, exe, train_info_dict, eval_info_dict)
2020-09-01 13:44:51 +08:00
elif train_alg_type == 'rec':
2020-05-10 16:26:57 +08:00
program.train_eval_rec_run(config, exe, train_info_dict, eval_info_dict)
2020-09-01 13:44:51 +08:00
else:
program.train_eval_cls_run(config, exe, train_info_dict, eval_info_dict)
2020-05-10 16:26:57 +08:00
def test_reader():
2020-07-16 20:14:46 +08:00
logger.info(config)
train_reader = reader_main(config=config, mode="train")
import time
starttime = time.time()
count = 0
try:
for data in train_reader():
count += 1
if count % 1 == 0:
batch_time = time.time() - starttime
starttime = time.time()
2020-11-21 22:11:07 +08:00
logger.info("[reader]count: {}, data length: {}, time: {}".
format(count, len(data), batch_time))
except Exception as e:
2020-07-01 13:09:44 +08:00
logger.info(e)
logger.info("finish reader: {}, Success!".format(count))
2020-05-10 16:26:57 +08:00
if __name__ == '__main__':
2020-11-11 18:41:22 +08:00
enable_static_mode()
2020-09-01 13:44:51 +08:00
startup_program, train_program, place, config, train_alg_type = program.preprocess(
)
2020-11-21 22:11:07 +08:00
# run the train process
2020-05-10 16:26:57 +08:00
main()
2020-11-21 22:11:07 +08:00
# if you want to check the reader, you can comment `main` and run test_reader
# test_reader()